1ST-PLACE-WINNER/MiniMax-H3

1# SPDX-License-Identifier: Apache-2.02# Torch-native normalization for the MiniMax H3 visual VAE.3import math4import os5 6import torch7import torch.distributed as dist8import torch.nn as nn9import torch.nn.functional as F10 11from .conv import SpatialParallelConv3d12from .parallel import all_reduce, get_parallel_state13 14 15def _validate_activation(activation):16    valid_activations = {"identity", "silu", "relu"}17    if activation not in valid_activations:18        raise ValueError(19            f"Unsupported activation: {activation}. Supported: {valid_activations}"20        )21 22 23def _apply_activation(x, activation):24    _validate_activation(activation)25    if activation == "identity":26        return x27    if activation == "silu":28        return F.silu(x)29    return F.relu(x)30 31 32def _merge_time_to_batch(x):33    batch, channels, depth, height, width = x.shape34    return (35        x.permute(0, 2, 1, 3, 4)36        .contiguous()37        .view(batch * depth, channels, 1, height, width)38    )39 40 41def _split_time_from_batch(x, batch):42    batch_depth, channels, _, height, width = x.shape43    depth = batch_depth // batch44    return (45        x.view(batch, depth, channels, height, width)46        .permute(0, 2, 1, 3, 4)47        .contiguous()48    )49 50 51def fused_group_norm(x, num_groups, weight, bias, eps=1e-5, activation="silu"):52    out = F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps)53    return _apply_activation(out, activation)54 55 56def fused_spatial_norm(57    f,58    num_groups,59    norm_weight,60    norm_bias,61    dynamic_scale,62    dynamic_bias,63    eps=1e-5,64    activation="silu",65):66    norm_f = F.group_norm(67        f,68        num_groups,69        weight=norm_weight,70        bias=norm_bias,71        eps=eps,72    )73    out = norm_f * dynamic_scale + dynamic_bias74    return _apply_activation(out, activation)75 76 77class DummyAffine(torch.nn.Module):78    def __init__(self, num_channels, affine=True):79        super().__init__()80        if affine:81            self.weight = torch.nn.Parameter(torch.ones(num_channels))82            self.bias = torch.nn.Parameter(torch.zeros(num_channels))83        else:84            self.register_parameter("weight", None)85            self.register_parameter("bias", None)86 87    def forward(self, input):88        if self.weight is None:89            return input90        shape = [1, -1] + [1] * (input.dim() - 2)91        return input * self.weight.view(*shape) + self.bias.view(*shape)92 93 94class FusedGroupNorm3D(torch.nn.Module):95    """Compatibility wrapper implemented with native PyTorch ops."""96 97    def __init__(98        self,99        num_groups,100        num_channels,101        eps=1e-5,102        affine=True,103        activation="silu",104        cond_channels=None,105        use_t_isolated_gn=False,106        padding_mode="zeros",107        padding_mode_t=None,108        causal=True,109    ):110        super().__init__()111        _validate_activation(activation)112        self.num_groups = num_groups113        self.num_channels = num_channels114        self.eps = eps115        self.affine = affine116        self.activation = activation117        self.use_t_isolated_gn = use_t_isolated_gn118 119        if cond_channels is not None:120            self.use_spatial_affine = True121            self.norm_layer = DummyAffine(num_channels, affine=affine)122            self.conv_y = SpatialParallelConv3d(123                cond_channels,124                num_channels,125                kernel_size=1,126                padding_mode=padding_mode,127                padding_mode_t=padding_mode_t,128                causal=causal,129            )130            self.conv_b = SpatialParallelConv3d(131                cond_channels,132                num_channels,133                kernel_size=1,134                padding_mode=padding_mode,135                padding_mode_t=padding_mode_t,136                causal=causal,137            )138        else:139            self.use_spatial_affine = False140            if self.affine:141                self.weight = torch.nn.Parameter(torch.ones(num_channels))142                self.bias = torch.nn.Parameter(torch.zeros(num_channels))143            else:144                self.register_parameter("weight", None)145                self.register_parameter("bias", None)146 147    def forward(self, f, cond=None):148        need_reshape = self.use_t_isolated_gn and f.dim() == 5149        batch = f.shape[0] if need_reshape else None150        f_size = f.shape[-3:]151        if need_reshape:152            f = _merge_time_to_batch(f)153 154        if self.use_spatial_affine:155            scale = self.conv_y(cond)156            bias = self.conv_b(cond)157            if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:158                scale = F.interpolate(scale, size=f_size, mode="nearest")159                bias = F.interpolate(bias, size=f_size, mode="nearest")160            if need_reshape:161                scale = _merge_time_to_batch(scale)162                bias = _merge_time_to_batch(bias)163            out = fused_spatial_norm(164                f,165                self.num_groups,166                self.norm_layer.weight,167                self.norm_layer.bias,168                scale,169                bias,170                self.eps,171                self.activation,172            )173        else:174            if cond is not None:175                raise NotImplementedError("Dynamic affine is not defined")176            weight = self.weight if self.affine else None177            bias = self.bias if self.affine else None178            out = fused_group_norm(179                f, self.num_groups, weight, bias, self.eps, self.activation180            )181 182        if need_reshape:183            out = _split_time_from_batch(out, batch)184        return out185 186 187class SpatialParallelGroupNorm(nn.GroupNorm):188    def __init__(189        self,190        *args,191        **kwargs,192    ):193        super().__init__(*args, **kwargs)194        self.spatial_parallel = False195 196    def _compute_stats(self, input):197        batch, channels = input.shape[0], input.shape[1]198        spatial_dims = input.shape[2:]199        spatial_size = math.prod(spatial_dims)200 201        groups = self.num_groups202        x = input.reshape(batch, groups, channels // groups, -1).to(torch.float32)203 204        local_sum = x.sum(dim=(2, 3))205        local_square_sum = (x * x).sum(dim=(2, 3))206        local_n = (channels // groups) * spatial_size207        local_n_tensor = torch.full_like(local_sum, float(local_n))208 209        stats = torch.stack([local_sum, local_square_sum, local_n_tensor], dim=0)210 211        local_process_group = get_parallel_state()["local_process_group"]212        stats = all_reduce(stats, dist.ReduceOp.SUM, local_process_group)213 214        total_sum = stats[0]215        total_square_sum = stats[1]216        total_n = stats[2]217 218        mean = total_sum / total_n219        var = (total_square_sum / total_n) - mean**2220        return mean, var221 222    def forward(self, input):223        if not self.spatial_parallel:224            return nn.GroupNorm.forward(self, input)225 226        batch, channels = input.shape[0], input.shape[1]227        orig_shape = input.shape228 229        mean, var = self._compute_stats(input)230        x = input.reshape(batch, self.num_groups, channels // self.num_groups, -1)231 232        mean = mean.unsqueeze(-1).unsqueeze(-1)233        var = var.unsqueeze(-1).unsqueeze(-1)234        x = (x - mean) / torch.sqrt(var + self.eps)235        x = x.reshape(orig_shape)236 237        if self.affine:238            shape = [1, -1] + [1] * (len(orig_shape) - 2)239            x *= self.weight.view(*shape)240            x += self.bias.view(*shape)241 242        return x243 244 245class TemporalIsolatedSpatialParallelGroupNorm(SpatialParallelGroupNorm):246    def forward(self, input):247        if input.dim() == 5:248            batch = input.shape[0]249            input = _merge_time_to_batch(input)250            output = super().forward(input)251            return _split_time_from_batch(output, batch)252        return super().forward(input)253 254 255 256 257 258 259 260 261class SpatialNorm3D(nn.Module):262    def __init__(263        self,264        f_channels,265        zq_channels,266        padding_mode="zeros",267        padding_mode_t=None,268        causal=True,269        use_t_isolated_gn=False,270    ):271        super().__init__()272        norm_cls = (273            TemporalIsolatedSpatialParallelGroupNorm274            if use_t_isolated_gn275            else SpatialParallelGroupNorm276        )277        self.norm_layer = norm_cls(278            num_groups=32, num_channels=f_channels, eps=1e-6, affine=True279        )280 281        self.conv_y = SpatialParallelConv3d(282            zq_channels,283            f_channels,284            kernel_size=1,285            padding_mode=padding_mode,286            padding_mode_t=padding_mode_t,287            causal=causal,288        )289        self.conv_b = SpatialParallelConv3d(290            zq_channels,291            f_channels,292            kernel_size=1,293            padding_mode=padding_mode,294            padding_mode_t=padding_mode_t,295            causal=causal,296        )297 298    def forward(self, f, zq):299        f_size = f.shape[-3:]300        norm_f = self.norm_layer(f)301        scale = self.conv_y(zq)302        bias = self.conv_b(zq)303 304        if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:305            scale = F.interpolate(scale, size=f_size, mode="nearest")306            bias = F.interpolate(bias, size=f_size, mode="nearest")307 308        return norm_f * scale + bias309 310 311def get_spatial_norm_3d(312    num_channels,313    cond_channels,314    *,315    padding_mode="zeros",316    padding_mode_t=None,317    causal=True,318    use_t_isolated_gn=False,319):320    if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":321        return FusedGroupNorm3D(322            num_groups=32,323            num_channels=num_channels,324            eps=1e-6,325            affine=True,326            cond_channels=cond_channels,327            use_t_isolated_gn=use_t_isolated_gn,328            padding_mode=padding_mode,329            padding_mode_t=padding_mode_t,330            causal=causal,331        )332    return SpatialNorm3D(333        num_channels,334        cond_channels,335        padding_mode=padding_mode,336        padding_mode_t=padding_mode_t,337        causal=causal,338        use_t_isolated_gn=use_t_isolated_gn,339    )340 341 342def get_group_norm_3d(num_channels, use_t_isolated_gn=False):343    if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":344        return FusedGroupNorm3D(345            num_groups=32,346            num_channels=num_channels,347            eps=1e-6,348            affine=True,349            use_t_isolated_gn=use_t_isolated_gn,350        )351 352    norm_cls = (353        TemporalIsolatedSpatialParallelGroupNorm354        if use_t_isolated_gn355        else SpatialParallelGroupNorm356    )357    return norm_cls(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)358