Track upstream torchtitan changes up to 20260713

[hann-wang]

The major changes in upstream torchtitan is that mixed-precision is now handled by fsdp instead of autocast. This means that we need to deal with DTensor now even with one GPU only.

Replaces #10.

[Copilot]

Pull request overview

This PR updates ALTO to track upstream torchtitan changes through 2026-07-13, especially the shift to FSDP-managed mixed precision (and the resulting need to handle DTensor even on single GPU), plus associated config/script updates.

Changes:

• Added/adjusted DTensor handling in observers and sparsification (notably WANDA) to work with upstream FSDP mixed-precision behavior.
• Updated modifier lifecycle to support config-time conversion (convert_config / on_convert_config) and refreshed several modifier/model config patterns (e.g., ignore lm_head).
• Updated example/integration scripts and model registries for upstream API changes; bumped dev version and changelog.

Reviewed changes

Copilot reviewed 58 out of 58 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
version.txt	Bumps project dev version to 0.1.0-dev0.
CHANGELOG.md	Adds v0.1.0 dev notes (FSDP mixed precision, Instella config removal).
tests/integration/instella_3b_opt.sh	Removes Instella 3B integration script.
tests/integration/gpt_oss_debugmodel_lpt.sh	Adds expert parallel degree flag in integration run.
examples/run.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/WandaStructured_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/Wanda_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/SparseGPT_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/RTN_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/OBS_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/Magnitude_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/llama3.2_1b_mx9.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path (placement adjusted).
examples/GPTQ_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/CosineSimilarity_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/AWQ_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/ALPS_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/AdmmStructured_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
examples/ADMM_Llama_3.1_8B.sh	Sets HSA_NO_SCRATCH_RECLAIM=1 for torchrun path.
alto/utils/exportation/export.py	Removes tied-embedding hotfix during HF export conversion.
alto/train.py	Updates forward path to align with upstream PP schedule API and FSDP MP expectations.
alto/observers/per_channel_norm.py	Introduces DTensor-aware stats initialization and local handling.
alto/nn/decomposed_linear.py	Converts DecomposedLinear to torchtitan Module + nested Config.
alto/modifiers/sparsification/wanda.py	Adds DTensor support in WANDA sparsification (redistribute/to_local/from_local).
alto/modifiers/sparsification/base.py	Adds on_convert_config hook to sparsification base lifecycle.
alto/modifiers/quantization/base.py	Moves quantization init/validation into on_convert; adds config-convert hook.
alto/modifiers/pruning/base.py	Adds on_convert_config hook to pruning base lifecycle.
alto/modifiers/lpt/base.py	Updates SDPA wrapper integration and adds config-time conversion for Linear→DecomposedLinear.
alto/modifiers/distillation/base.py	Updates optimizer configuration to param-group style; adds config-convert hook.
alto/modifiers/base.py	Introduces convert_config + abstract on_convert_config for modifiers.
alto/models/patcher.py	Updates rotary-embedding patching to match upstream ComplexRoPE API.
alto/models/llama3/configs/wanda_structured_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/wanda_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/sparsegpt_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/rtn_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/recipe.yaml	Updates commented and active ignore list references to lm_head.
alto/models/llama3/configs/obs_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/mx9_wa_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/mx6_wa_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/magnitude_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/lpt_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/guanchen_recipe.yaml	Updates ignore list references to lm_head (including commented blocks).
alto/models/llama3/configs/gptq_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/cosine_similarity_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/awq_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/awq_debug_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/alps_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/admm_structured_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/configs/admm_recipe.yaml	Switches ignore list from output to lm_head.
alto/models/llama3/config_registry.py	Updates trainer config fields for upstream API changes (profiler, optimizer factory, activation checkpointing).
alto/models/gpt_oss/configs/lpt_recipe.yaml	Updates ignore list to lm_head and keeps router gate ignore.
alto/models/gpt_oss/config_registry.py	Updates profiler + activation checkpoint fields for upstream API changes.
alto/models/deepseek_v3/configs/lpt_recipe.yaml	Updates ignore list to lm_head and keeps router gate ignore.
alto/models/deepseek_v3/config_registry.py	Updates profiler + activation checkpoint fields; removes selective AC option.
alto/kernels/dispatch/conversion.py	Adjusts bias-parameter skip condition while swapping wrapped parameters.
alto/kernels/dispatch/attention.py	Replaces SDPA wrapper with LPScaledDotProductAttention matching upstream SDPA class/signature.
alto/kernels/dispatch/init.py	Updates exports/imports for renamed SDPA wrapper class.
alto/components/converter.py	Adds convert_config pass to apply modifier config conversions.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 59 out of 59 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 59 out of 59 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 59 out of 59 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 61 out of 61 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

CHANGELOG.md:8

• The changelog entry says “Dropped Instella-3B model config.”, but the Instella configs are still present and exported (e.g. alto/models/llama3/config_registry.py still defines/exports instella_3b*). The only removal in this PR is the Instella integration script, so the changelog bullet should be adjusted to avoid misleading users.

- Changed
  - Mixed precision is now handled by FSDP even if world_size=1.
  - Dropped Instella-3B model config.

tests/unittest/nn/test_decomposed_linear.py:165

• test_decomposed_linear_with_svd currently never asserts correctness—so it will pass even if outputs/gradients are wrong. Add basic thresholds on the computed SNR / cosine-sim metrics (similar to the quantization test above) so the test actually guards behavior.

[Copilot]

Pull request overview

Copilot reviewed 69 out of 69 changed files in this pull request and generated no new comments.

Suppressed comments (2)

tests/unittest/nn/test_decomposed_linear.py:15

• from utils import ... will fail under pytest because there is no top-level utils.py on the import path (and tests/ is not a package). Other test suites import these helpers from alto.kernels.fp4.testing_utils and/or define prepare_data locally (e.g., tests/unittest/mxfp4/utils.py).
  alto/modifiers/lpt/base.py:117
• Using assert inside a Pydantic validator makes the validation behavior dependent on Python optimization flags (python -O disables asserts). These should be normal validation errors so misconfigurations are always rejected.

[Copilot]

Pull request overview

Copilot reviewed 70 out of 70 changed files in this pull request and generated no new comments.

Suppressed comments (2)

tests/unittest/nn/test_decomposed_linear.py:14

• from utils import ... will fail because there is no top-level utils module in the repo. The helper functions live in tests/unittest/mxfp4/utils.py, so this import will raise ModuleNotFoundError when running this test.
  alto/modifiers/lpt/base.py:117
• The use_uos validation is currently after an early return (if self.lora_rank <= 0: return self), so when LoRA is disabled these constraints are never enforced. Also, assert inside pydantic validators can be optimized away with python -O; use explicit ValueErrors instead.

[Copilot]

Pull request overview

Copilot reviewed 71 out of 71 changed files in this pull request and generated no new comments.

Suppressed comments (5)

tests/unittest/nn/test_decomposed_linear.py:14

• from utils import ... will fail at import time because there is no top-level utils.py on the test import path for this directory. The shared helpers live under the mxfp4 test package (and already provide prepare_data, calc_snr, calc_cossim).
  alto/modifiers/lpt/base.py:42
• The default ignore list for this modifier is still ["output"] (see class fields above), but configs/recipes in this PR have switched to excluding lm_head. With the current default, lm_head may be unintentionally converted unless every caller overrides ignore. Consider updating the default to lm_head for consistency with the updated recipes.
  tests/unittest/mxfp4/utils.py:332
• In the use_uos path, scales can exceed 255 for large-magnitude blocks; casting to uint8 later will wrap and produce incorrect scaling factors. Clamp to 255 before the cast (similar issue exists in the Triton kernel implementation).
  alto/modifiers/lpt/base.py:117
• Using assert inside a Pydantic @model_validator is unsafe because assertions can be stripped with python -O, disabling this validation. Prefer raising ValueError so invalid configs are always rejected.
  alto/kernels/fp4/mxfp4/mxfp_quantization.py:78
• In the USE_UOS path, scales can exceed 255 for large-magnitude blocks; later casting to tl.uint8 will wrap and yield incorrect scales. Clamp the upper bound before converting to tl.uint8.

    if USE_UOS:
        scales = tl.ceil(tl.log2(max_abs / 7.25))
        scales = tl.where(max_abs == 0, 1.0, scales)
        scales += 127

[Copilot]

Pull request overview

Copilot reviewed 71 out of 71 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

alto/modifiers/lpt/base.py:117

• This pydantic validator uses assert for user-config validation. assert statements can be stripped with Python optimizations (-O), turning these into silent misconfigurations. Prefer raising ValueError so invalid configs always fail fast.

[Copilot]

Pull request overview

Copilot reviewed 71 out of 71 changed files in this pull request and generated 2 comments.

Suppressed comments (3)

CHANGELOG.md:7

• Changelog says "Dropped Instella-3B model config", but Instella configs (instella_3b*) are still present in alto/models/llama3/config_registry.py. In this PR the only Instella-related removal appears to be the tests/integration/instella_3b_opt.sh script, so the changelog entry is currently inaccurate.

  - Mixed precision is now handled by FSDP even if world_size=1.
  - Dropped Instella-3B model config.

alto/modifiers/distillation/base.py:184

• global_valid_tokens passed into forward_step() is currently cumulative across all prior microbatch groups in this loop (local_valid_tokens is never reset). This changes semantics vs. the previous per-microbatch token count and can skew loss normalization/scaling inside forward_step(). Compute per-group valid tokens instead.
  alto/modifiers/lpt/base.py:117
• Using assert inside a Pydantic @model_validator is unsafe because assertions can be stripped with python -O, turning these validations into no-ops. Raise ValueError for invalid configurations instead.

[Copilot]

Pull request overview

Copilot reviewed 58 out of 58 changed files in this pull request and generated no new comments.

Suppressed comments (4)

alto/train.py:191

• In PP mode, forward_step() calls self.pp_forward_backward_step(...), but no such method exists in this file/repo. This will raise AttributeError at runtime. It looks like the intended call is the newly added pp_forward_step().
  alto/train.py:266
• pp_forward_step() returns None on non-last pipeline stages, but train_step() unconditionally does result.detach().cpu() and caches it. On non-last PP stages this will crash with 'NoneType' has no attribute 'detach'. Either only cache outputs on the last stage or guard the detach/caching when result is None.
  alto/modifiers/lpt/base.py:118
• validate_lora_rank_alignment() uses assert for user configuration validation (use_uos constraints). assert can be stripped with Python optimizations (-O) and will raise AssertionError instead of a Pydantic validation error. Prefer raising ValueError for invalid configs.
  alto/kernels/fp4/mxfp4/mxfp_quantization.py:558
• The @convert_to_mxfp4.register_fake implementation has a different signature than the real convert_to_mxfp4 op (missing philox_seed, philox_offset, and clip_mode, and instead taking use_static_clip). In FakeTensor/torch.compile/meta contexts, this can raise a TypeError when the op is invoked with the real kwargs.

@convert_to_mxfp4.register_fake
def _fake_convert_to_mxfp4(
    data_hp: torch.Tensor,
    block_size: int = BLOCK_SIZE_DEFAULT,
    axis: int = -1,
    is_2d_block: bool = False,
    use_sr: bool = False,
    use_asm: Optional[bool] = None,
    use_static_clip: bool = False,
    use_uos: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]: