Track upstream torchtitan changes up to 20260713#35
Conversation
There was a problem hiding this comment.
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., ignorelm_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.
| new_config = DecomposedLinear.Config( | ||
| in_features=config.in_features, | ||
| out_features=config.out_features, | ||
| bias=config.bias, | ||
| param_init=config.param_init | DecomposedLinear._EXTRA_INIT, | ||
| ) | ||
| setattr(parent, attr, new_config) |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
| new_config = DecomposedLinear.Config( | ||
| in_features=config.in_features, | ||
| out_features=config.out_features, | ||
| bias=config.bias, | ||
| param_init=config.param_init | DecomposedLinear._EXTRA_INIT, | ||
| ) |
| - Changed | ||
| - Mixed precision is now handled by FSDP even if world_size=1. | ||
| - Dropped Instella-3B model config. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 60 out of 60 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
alto/modifiers/lpt/base.py:225
- In on_convert_config(),
lora_rankis pulled fromconfig.lora_rank, butconfighere is aLinear.Configproduced by torchtitan. This is inconsistent with the runtime conversion path (which usesself.lora_rank) and will likely raise an AttributeError or silently use the wrong rank. Use the modifier-levelself.lora_rankwhen constructingDecomposedLinear.Config.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 60 out of 60 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
alto/modifiers/quantization/base.py:126
- on_convert() unconditionally writes to mod._param_init for modules that have quantization_scheme. Target modules are often plain torch.nn.Modules (e.g., torch.nn.Linear) which do not define _param_init, so this will raise AttributeError during conversion.
| def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: | ||
| if not QuantizationMixin.has_config(self): | ||
| raise ValueError("QuantizationModifier requires that quantization fields be specified") | ||
| for m in model_parts: | ||
| QuantizationMixin.initialize_quantization(self, m) | ||
|
|
||
| if self.sequential: | ||
| self._build_sequential_blocks(model_parts) | ||
| return True |
| *, | ||
| attention_masks: None = None, | ||
| scale: float | None = None, |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 60 out of 60 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
CHANGELOG.md:7
- CHANGELOG entry says "Dropped Instella-3B model config", but Instella configs (instella_3b/instella_3b_opt/instella_3b_lpt) are still present in alto/models/llama3/config_registry.py. Either remove those configs or update the changelog to reflect what was actually removed (e.g., only the integration script).
- Mixed precision is now handled by FSDP even if world_size=1.
- Dropped Instella-3B model config.
| def on_convert(self, model: Module, **kwargs) -> bool: | ||
| for scheme_obj, targets in self.resolved_config.items(): | ||
| for name, module in match_named_modules(model, targets, self.ignore): | ||
| if isinstance(module, BaseAttention): |
| if self.lora_rank > 0: | ||
| resolved_targets = list(self.resolved_config.values()) | ||
| for _fqn, config, parent, attr in model_config.traverse(Linear.Config): | ||
| for target in resolved_targets: | ||
| if is_match(_fqn, "Linear", target, self.ignore): |
The major changes in upstream torchtitan is that mixed-precision is now handled by
fsdpinstead ofautocast. This means that we need to deal withDTensornow even with one GPU only.Replaces #10.