Track upstream torchtitan changes up to 20260713

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## v0.1.0 [dev]
+
+- Changed
+  - Mixed precision is now handled by FSDP even if world_size=1.
+  - Dropped Instella-3B model config.
+
+
 ## v0.0.2 [dev]
 
 - Changed

alto/components/converter.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
+from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
 import torch
@@ -15,6 +16,9 @@
 
 from alto.config import Recipe
 
+if TYPE_CHECKING:
+    from torchtitan.models.base import BaseModel
+
 
 class ModelOptConverter(ModelConverter, Configurable):
 
@@ -45,6 +49,10 @@ def convert(self, model: nn.Module):
 
         for modifier in self.recipe.modifiers:
             modifier.convert(model)
+
+    def convert_config(self, model_config: "BaseModel.Config"):
+        for modifier in self.recipe.modifiers:
+            modifier.convert_config(model_config)
 
     def pre_step(self, model_parts: list[nn.Module], **kwargs):
         for modifier in self.recipe.modifiers:

alto/components/optimizer.py

@@ -153,7 +153,7 @@ def qdq(w: torch.Tensor, axis: int) -> torch.Tensor:
                 w = torch.nn.functional.pad(w, (0, 0, 0, 32 - original_rows % 32))
 
             data_lp, scales = convert_to_mxfp4(
-                w, axis=axis, is_2d_block=is_2d_block,
+                w, axis=axis, is_2d_block=is_2d_block, use_uos=cfg.use_uos
             )
             dequantized = convert_from_mxfp4(
                 data_lp,

alto/kernels/dispatch/__init__.py

@@ -4,10 +4,10 @@
 
 from .config import TrainingOpConfig
 from .conversion import swap_params
-from .attention import LPScaledDotProductAttentionWrapper
+from .attention import LPScaledDotProductAttention
 
 __all__ = [
     "TrainingOpConfig",
     "swap_params",
-    "LPScaledDotProductAttentionWrapper",
+    "LPScaledDotProductAttention",
 ]

alto/kernels/dispatch/attention.py

@@ -3,15 +3,15 @@
 # SPDX-License-Identifier: MIT
 
 import torch
-from torchtitan.models.common.attention import (ScaledDotProductAttentionWrapper)
+from torchtitan.models.common.attention import (ScaledDotProductAttention)
 
 from alto.kernels.fp4.mxfp4.triton_flash_attention_mxfp4 import triton_attention_mxfp4
 from .config import TrainingOpConfig
 
-__all__ = ["LPScaledDotProductAttentionWrapper"]
+__all__ = ["LPScaledDotProductAttention"]
 
 
-class LPScaledDotProductAttentionWrapper(ScaledDotProductAttentionWrapper):
+class LPScaledDotProductAttention(ScaledDotProductAttention):
 
     def __init__(self, config: TrainingOpConfig):
         super().__init__()
@@ -25,17 +25,30 @@ def __init__(self, config: TrainingOpConfig):
 
     def _get_name(self) -> str:
         return f"{self.__class__.__name__}[{self.config}]"
-
+    
     def forward(
         self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
+        q_BLNH: torch.Tensor,
+        k_BLNH: torch.Tensor,
+        v_BLNH: torch.Tensor,
         *,
+        attention_masks: None = None,
         scale: float | None = None,
         enable_gqa: bool = False,
         is_causal: bool = True,
-    ):
+        **kwargs,
+    ) -> torch.Tensor:
+        if attention_masks is not None:
+            raise ValueError(
+                "ScaledDotProductAttention does not support attention_masks; it "
+                "only supports causal/non-causal attention via is_causal."
+            )
+        # Transpose to (B, N, L, H) for SDPA
+        q, k, v = (
+            q_BLNH.transpose(1, 2),
+            k_BLNH.transpose(1, 2),
+            v_BLNH.transpose(1, 2),
+        )
         batch, num_head_q, seqlen_q, head_dim_qk = q.shape
         batch_k, num_head_kv, seqlen_kv, head_dim_qk_k = k.shape
         batch_v, num_head_kv_v, seqlen_kv_v, head_dim_v = v.shape
@@ -64,4 +77,5 @@ def forward(
             use_exp2=True,
             layout="bhsd",
         )[0]
-        return o
+        # Transpose back to (B, L, N, H)
+        return o.transpose(1, 2)

alto/kernels/dispatch/config.py

@@ -47,6 +47,11 @@ class TrainingOpConfig:
       * MXFP4: apply a blockwise scale factor containing shared mantissa to each block
       * NVFP4: not implemented
     """
+
+    use_uos: bool = False
+    """
+    use 7.25 instead of 6 for MXFP4 target range.
+    """
 
 
 torch.serialization.add_safe_globals([TrainingOpConfig])

alto/kernels/dispatch/conversion.py

@@ -101,7 +101,7 @@ def post_order_traversal(
                     continue
                 module_prefix = f"{module_name}." if module_name else ""
                 full_param_name = f"{module_prefix}{cur_fqn}{'.' if cur_fqn else ''}{param_name}"
-                if target_parameter_name is None and param_name.endswith("bias"):
+                if target_parameter_name is None and "bias" in param_name:
                     logger.debug(f"Skipped {full_param_name} because it is a bias parameter")
                     continue
                 if not isinstance(param.data, TrainingWeightWrapperBaseTensor):

alto/kernels/dispatch/tensor.py

@@ -152,6 +152,9 @@ def __tensor_unflatten__(cls, inner_tensors, flatten_spec, outer_size, outer_str
             flatten_spec["config"],
         )
 
+    def untyped_storage(self):
+        return self._data.untyped_storage()
+
     # fsdp hooks based on https://github.com/pytorch/pytorch/blob/20e40492b046b9287726d3ec656117e4dc38f0e2/test/distributed/_composable/fsdp/test_fully_shard_extensions.py#L81
     def fsdp_pre_all_gather(
         self,
@@ -256,6 +259,7 @@ def __torch_function__(cls, func, types, args, kwargs={}):
                 use_hadamard=config.use_hadamard,
                 clip_mode=config.clip_mode,
                 use_macro_block_scaling=config.two_level_scaling == "blockwise",
+                use_uos=config.use_uos,
             )
 
         # linear op override
@@ -287,6 +291,7 @@ def __torch_function__(cls, func, types, args, kwargs={}):
                 clip_mode=config.clip_mode,
                 use_hadamard=config.use_hadamard,
                 use_macro_block_scaling=config.two_level_scaling == "blockwise",
+                use_uos=config.use_uos,
             )
             if bias is not None:
                 Y = Y + bias

alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/cg_backward.py

@@ -616,6 +616,7 @@ def forward(
         clip_mode=False,
         use_macro_block_scaling=False,
         hadamard_transform: Optional[HadamardTransform] = None,
+        use_uos: bool = False,
     ):
         """Forward pass for contiguous grouped GEMM."""
         original_dtype = inputs.dtype
@@ -641,11 +642,13 @@ def forward(
             inputs_scaled,
             axis=-1,
             is_2d_block=use_2dblock_x,
+            use_uos=use_uos,
         )
         expert_weights_mxfp4, expert_weight_scales = torch.ops.torchtitan.convert_to_mxfp4(
             expert_weights_scaled,
             axis=quant_axis_w,
             is_2d_block=use_2dblock_w,
+            use_uos=use_uos,
         )
 
         if is_cdna4():
@@ -697,6 +700,7 @@ def forward(
                 expert_weights_scaled,
                 axis=requant_axis_w,
                 is_2d_block=False,
+                use_uos=use_uos,
             )
             if not is_cdna4():
                 w_dq = torch.ops.torchtitan.convert_from_mxfp4(
@@ -724,6 +728,7 @@ def forward(
                 axis=0,
                 is_2d_block=False,
                 clip_mode=clip_mode,
+                use_uos=use_uos,
             )
             if not is_cdna4():
                 x_dq = torch.ops.torchtitan.convert_from_mxfp4(
@@ -752,6 +757,7 @@ def forward(
         ctx.hadamard_transform = hadamard_transform
         ctx.clip_mode = clip_mode
         ctx.use_macro_block_scaling = use_macro_block_scaling
+        ctx.use_uos = use_uos
 
         return res
 
@@ -785,6 +791,7 @@ def backward(ctx, grad_output):
                 axis=-1,
                 use_sr=ctx.use_sr_grad,
                 is_2d_block=True,
+                use_uos=ctx.use_uos,
             )
             grad_output_mxfp4_m = grad_output_mxfp4
             grad_output_scales_m = grad_output_scales
@@ -812,6 +819,7 @@ def backward(ctx, grad_output):
                 axis=-1,
                 use_sr=ctx.use_sr_grad,
                 is_2d_block=False,
+                use_uos=ctx.use_uos,
             )
             if ctx.hadamard_transform is not None:
                 grad_output = ctx.hadamard_transform(grad_output, left_mul=True)
@@ -826,6 +834,7 @@ def backward(ctx, grad_output):
                 use_sr=ctx.use_sr_grad,
                 is_2d_block=False,
                 clip_mode=ctx.clip_mode,
+                use_uos=ctx.use_uos,
             )
 
             if not is_cdna4():
@@ -918,7 +927,7 @@ def backward(ctx, grad_output):
             )
             grad_weights *= dge_bwd(w_fp4_values, torch.float4_e2m1fn_x2)
 
-        return grad_inputs, grad_weights, None, None, None, None, None, None, None, None, None
+        return grad_inputs, grad_weights, None, None, None, None, None, None, None, None, None, None
 
 
 def mxfp4_grouped_gemm(
@@ -934,6 +943,7 @@ def mxfp4_grouped_gemm(
     use_hadamard: bool = False,
     clip_mode: str = "none",
     use_macro_block_scaling: bool = False,
+    use_uos: bool = False
 ) -> torch.Tensor:
     """
     Interface for contiguous grouped GEMM with full backward pass support.
@@ -975,6 +985,7 @@ def mxfp4_grouped_gemm(
         clip_mode,
         use_macro_block_scaling,
         hadamard_transform,
+        use_uos,
     )
 
     return res

alto/kernels/fp4/mxfp4/mxfp_grouped_gemm/functional.py

@@ -19,6 +19,7 @@ def _quantize_then_mxfp_scaled_grouped_mm(
     use_hadamard: bool,
     clip_mode: str,
     use_macro_block_scaling: bool,
+    use_uos: bool = False,
 ) -> torch.Tensor:
     m_indices = create_indices_from_offsets_nosync(offs)
     return mxfp4_grouped_gemm(
@@ -33,4 +34,5 @@ def _quantize_then_mxfp_scaled_grouped_mm(
         use_hadamard=use_hadamard,
         clip_mode=clip_mode,
         use_macro_block_scaling=use_macro_block_scaling,
+        use_uos=use_uos,
     )

alto/kernels/fp4/mxfp4/mxfp_linear.py

@@ -274,6 +274,7 @@ def forward(
         clip_mode,
         use_macro_block_scaling,
         hadamard_transform: Optional[HadamardTransform] = None,
+        use_uos: bool = False,
     ):
         """
         Forward pass for the blockwise FP8 linear operation.
@@ -310,12 +311,14 @@ def forward(
             x_scaled,
             axis=-1,
             is_2d_block=use_2dblock_x,
+            use_uos=use_uos,
         )
 
         w_mxfp4, w_scale = torch.ops.torchtitan.convert_to_mxfp4(
             w_scaled,
             axis=-1,
             is_2d_block=use_2dblock_w,
+            use_uos=use_uos,
         )
 
         if is_cdna4():
@@ -363,6 +366,7 @@ def forward(
                 w_scaled,
                 axis=0,
                 is_2d_block=False,
+                use_uos=use_uos,
             )
             if not is_cdna4():
                 w_dq = torch.ops.torchtitan.convert_from_mxfp4(
@@ -388,6 +392,7 @@ def forward(
                 axis=0,
                 is_2d_block=False,
                 clip_mode=clip_mode,
+                use_uos=use_uos,
             )
             if not is_cdna4():
                 x_dq = torch.ops.torchtitan.convert_from_mxfp4(
@@ -412,6 +417,7 @@ def forward(
         ctx.use_dge = use_dge
         ctx.clip_mode = clip_mode
         ctx.use_macro_block_scaling = use_macro_block_scaling
+        ctx.use_uos = use_uos
 
         return y.view(*original_shape[:-1], -1)  # Reshape back to original
 
@@ -440,6 +446,7 @@ def backward(ctx, grad_output):
                 axis=-1,
                 is_2d_block=True,
                 use_sr=ctx.use_sr_grad,
+                use_uos=ctx.use_uos,
             )
             grad_output_mxfp4_m = grad_output_mxfp4
             grad_output_scales_m = grad_output_scales
@@ -467,6 +474,7 @@ def backward(ctx, grad_output):
                 axis=-1,
                 use_sr=ctx.use_sr_grad,
                 is_2d_block=False,
+                use_uos=ctx.use_uos,
             )
 
             if ctx.hadamard_transform is not None:
@@ -482,6 +490,7 @@ def backward(ctx, grad_output):
                 use_sr=ctx.use_sr_grad,
                 is_2d_block=False,
                 clip_mode=ctx.clip_mode,
+                use_uos=ctx.use_uos,
             )
 
             if not is_cdna4():
@@ -558,7 +567,7 @@ def backward(ctx, grad_output):
             )
             grad_weights *= dge_bwd(w_fp4_values, torch.float4_e2m1fn_x2)
 
-        return grad_inputs.view(*original_shape[:-1], -1), grad_weights, None, None, None, None, None, None, None
+        return grad_inputs.view(*original_shape[:-1], -1), grad_weights, None, None, None, None, None, None, None, None
 
 
 def _to_mxfp4_then_scaled_mm(
@@ -571,6 +580,7 @@ def _to_mxfp4_then_scaled_mm(
     clip_mode: str,
     use_hadamard: bool,
     use_macro_block_scaling: bool = False,
+    use_uos: bool = False,
 ) -> torch.Tensor:
     if use_hadamard:
         with torch.no_grad():
@@ -587,5 +597,6 @@ def _to_mxfp4_then_scaled_mm(
         clip_mode,
         use_macro_block_scaling,
         hadamard_transform,
+        use_uos,
     )
     return y