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/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/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,

alto/models/deepseek_v3/config_registry.py

@@ -21,12 +21,12 @@
 
 def deepseek_v3_debugmodel() -> Trainer.Config:
     config = deepseek_v3_debugmodel_orig()
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 10
     config.training.local_batch_size = 4
     config.training.global_batch_size = 16
     config.training.seq_len = 2048
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -43,15 +43,14 @@ def deepseek_v3_16b() -> Trainer.Config:
     config = deepseek_v3_16b_orig()
     config.hf_assets_path = "/huggingface/hub/models--deepseek-ai--deepseek-moe-16b-base/snapshots/521d2bc4fb69a3f3ae565310fcc3b65f97af2580"
     config.dump_folder = "deepseek_v3_16b-outputs"
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.seq_len = 4096
     config.metrics.log_freq = 1
     config.metrics.enable_tensorboard = True
     config.dataloader.dataset = "c4_test"
     config.parallelism.expert_parallel_degree = 1
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
     config.checkpoint.enable = True
     config.checkpoint.initial_load_path = "/huggingface/hub/models--deepseek-ai--deepseek-moe-16b-base/snapshots/521d2bc4fb69a3f3ae565310fcc3b65f97af2580"
@@ -61,8 +60,7 @@ def deepseek_v3_16b() -> Trainer.Config:
     config.validator.dataloader.dataset = "wikitext_test"
     config.validator.freq = 10
     config.validator.steps = 10
-    config.activation_checkpoint.mode = "none"
-    config.activation_checkpoint.selective_ac_option = "1"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 

alto/models/deepseek_v3/configs/lpt_recipe.yaml

@@ -3,7 +3,7 @@ training_stage:
     LowPrecisionTrainingModifier:
       scheme: "mxfp4"
       targets: ["Linear", "GroupedExperts"]
-      ignore: ["output", "re:.*\\.router\\.gate"]
+      ignore: ["lm_head", "re:.*\\.router\\.gate"]
       use_2dblock_x: false
       use_2dblock_w: true
       use_hadamard: true

alto/models/gpt_oss/config_registry.py

@@ -22,12 +22,12 @@
 
 def gpt_oss_debugmodel() -> Trainer.Config:
     config = gpt_oss_debugmodel_orig()
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 10
     config.training.local_batch_size = 4
     config.training.global_batch_size = 16
     config.training.seq_len = 2048
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -44,15 +44,14 @@ def gpt_oss_20b() -> Trainer.Config:
     config = gpt_oss_20b_orig()
     config.hf_assets_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/"
     config.dump_folder = "gpt_oss_20b-outputs"
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 0
     config.training.local_batch_size = 1
     config.training.seq_len = 8192
     config.metrics.log_freq = 1
     config.metrics.enable_tensorboard = True
     config.dataloader.dataset = "c4_test"
     config.parallelism.expert_parallel_degree = 1
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
     config.checkpoint.enable = True
     config.checkpoint.initial_load_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/"
@@ -63,7 +62,7 @@ def gpt_oss_20b() -> Trainer.Config:
     config.validator.dataloader.dataset = "wikitext_test"
     config.validator.freq = 10
     config.validator.steps = 10
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
@@ -72,16 +71,14 @@ def gpt_oss_20b_pretrain() -> Trainer.Config:
     config = gpt_oss_20b_orig()
     config.hf_assets_path = "/huggingface/hub/models--openai--gpt-oss-20b/snapshots/6cee5e81ee83917806bbde320786a8fb61efebee/"
     config.dump_folder = "gpt_oss_20b-pretrain-subset-lr4e-4-outputs"
-    config.profiling.enable_profiling = False
+    config.profiler.enable_profiling = False
     config.training.steps = 1200000
     config.training.local_batch_size = 1
     config.training.global_batch_size = 16
     config.training.seq_len = 8192
-    config.optimizer.lr = 4e-4
-    config.optimizer.weight_decay = 0.1
-    config.optimizer.beta1 = 0.9
-    config.optimizer.beta2 = 0.95
-    config.optimizer.eps = 1e-5
+    config.optimizer.param_groups[0].optimizer_kwargs["lr"] = 4e-4
+    config.optimizer.param_groups[0].optimizer_kwargs["betas"] = (0.9, 0.95)
+    config.optimizer.param_groups[0].optimizer_kwargs["eps"] = 1e-5
     config.lr_scheduler.min_lr_factor = 0.1
     config.lr_scheduler.warmup_steps = 128
     config.lr_scheduler.decay_ratio = 1 - 128 / config.training.steps
@@ -91,7 +88,6 @@ def gpt_oss_20b_pretrain() -> Trainer.Config:
     config.dataloader.dataset = "megatron"
     config.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-train.en_6_text_document.idx"
     config.parallelism.expert_parallel_degree = 8
-    config.parallelism.expert_tensor_parallel_degree = 1
     config.parallelism.tensor_parallel_degree = 1
     config.checkpoint.enable = True
     config.checkpoint.interval = 1000
@@ -101,14 +97,14 @@ def gpt_oss_20b_pretrain() -> Trainer.Config:
     config.validator.dataloader.dataset_path = "/workspace/workspace/megatron_dataset/data/c4-validation-91205-samples.en_text_document.idx"
     config.validator.freq = 768
     config.validator.steps = 64
-    config.activation_checkpoint.mode = "none"
+    config.activation_checkpoint = None
     config.debug.seed = 1234
     return config
 
 
 def gpt_oss_20b_lpt() -> Trainer.Config:
     config = gpt_oss_20b_pretrain()
-    config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-outputs"
+    config.dump_folder = "gpt_oss_20b-pretrain-subset-mxfp4gemm_1d2d-hadamard-sr-lr4e-4-refactor-outputs"
     config.model_converters = ModelConvertersContainer.Config(converters=[
         ModelOptConverter.Config(recipe="./alto/models/gpt_oss/configs/lpt_recipe.yaml",),
     ],)

alto/models/gpt_oss/configs/lpt_recipe.yaml

@@ -4,7 +4,7 @@ training_stage:
       scheme: "mxfp4"
       targets: ["Linear", "GptOssGroupedExperts"]
       # targets: ["Linear"]
-      ignore: ["output", "re:.*\\.router\\.gate"]
+      ignore: ["lm_head", "re:.*\\.router\\.gate"]
       use_2dblock_x: false
       use_2dblock_w: true
       use_hadamard: true