From 9f19b2cebea53421078ed2874c26c434beee6336 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 14 Jul 2026 03:48:22 +0000 Subject: [PATCH 01/14] feat: adapt torchtitan dev 20260713 --- 3rdparty/torchtitan | 2 +- alto/kernels/dispatch/__init__.py | 4 +- alto/kernels/dispatch/attention.py | 32 ++++++--- alto/kernels/dispatch/conversion.py | 2 +- alto/models/deepseek_v3/config_registry.py | 9 ++- .../deepseek_v3/configs/lpt_recipe.yaml | 2 +- alto/models/gpt_oss/config_registry.py | 12 ++-- alto/models/gpt_oss/configs/lpt_recipe.yaml | 2 +- alto/models/llama3/config_registry.py | 20 +++--- alto/models/llama3/configs/admm_recipe.yaml | 2 +- .../configs/admm_structured_recipe.yaml | 2 +- alto/models/llama3/configs/alps_recipe.yaml | 2 +- .../llama3/configs/awq_debug_recipe.yaml | 2 +- alto/models/llama3/configs/awq_recipe.yaml | 2 +- .../configs/cosine_similarity_recipe.yaml | 2 +- alto/models/llama3/configs/gptq_recipe.yaml | 2 +- .../llama3/configs/guanchen_recipe.yaml | 8 +-- alto/models/llama3/configs/lpt_recipe.yaml | 2 +- .../llama3/configs/magnitude_recipe.yaml | 2 +- alto/models/llama3/configs/mx6_wa_recipe.yaml | 2 +- alto/models/llama3/configs/mx9_wa_recipe.yaml | 2 +- alto/models/llama3/configs/obs_recipe.yaml | 2 +- alto/models/llama3/configs/recipe.yaml | 4 +- alto/models/llama3/configs/rtn_recipe.yaml | 2 +- .../llama3/configs/sparsegpt_recipe.yaml | 2 +- alto/models/llama3/configs/wanda_recipe.yaml | 2 +- .../configs/wanda_structured_recipe.yaml | 2 +- alto/models/patcher.py | 68 +++++++++---------- alto/modifiers/base.py | 10 +++ alto/modifiers/distillation/base.py | 7 +- alto/modifiers/lpt/base.py | 26 ++++--- alto/modifiers/pruning/base.py | 7 +- alto/modifiers/quantization/base.py | 8 ++- alto/modifiers/sparsification/base.py | 7 +- alto/train.py | 11 ++- alto/utils/exportation/export.py | 12 ---- examples/ADMM_Llama_3.1_8B.sh | 1 + examples/ALPS_Llama_3.1_8B.sh | 1 + examples/AWQ_Llama_3.1_8B.sh | 1 + examples/AdmmStructured_Llama_3.1_8B.sh | 1 + examples/CosineSimilarity_Llama_3.1_8B.sh | 1 + examples/GPTQ_Llama_3.1_8B.sh | 1 + examples/Magnitude_Llama_3.1_8B.sh | 1 + examples/OBS_Llama_3.1_8B.sh | 1 + examples/RTN_Llama_3.1_8B.sh | 1 + examples/SparseGPT_Llama_3.1_8B.sh | 1 + examples/WandaStructured_Llama_3.1_8B.sh | 1 + examples/Wanda_Llama_3.1_8B.sh | 1 + examples/llama3.2_1b_mx9.sh | 1 + examples/run.sh | 1 + 50 files changed, 180 insertions(+), 119 deletions(-) diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan index 82084e7c..8296b499 160000 --- a/3rdparty/torchtitan +++ b/3rdparty/torchtitan @@ -1 +1 @@ -Subproject commit 82084e7c67c895e1f91b6d90887d3f7c83909577 +Subproject commit 8296b49912ddd98b75a0584e15ce04ff0fc9a9db diff --git a/alto/kernels/dispatch/__init__.py b/alto/kernels/dispatch/__init__.py index 2e03babe..fdb6beac 100644 --- a/alto/kernels/dispatch/__init__.py +++ b/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", ] diff --git a/alto/kernels/dispatch/attention.py b/alto/kernels/dispatch/attention.py index 33727623..b77658f2 100644 --- a/alto/kernels/dispatch/attention.py +++ b/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) diff --git a/alto/kernels/dispatch/conversion.py b/alto/kernels/dispatch/conversion.py index 325c9e35..8b8e35b6 100644 --- a/alto/kernels/dispatch/conversion.py +++ b/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): diff --git a/alto/models/deepseek_v3/config_registry.py b/alto/models/deepseek_v3/config_registry.py index dc0f0ec9..14e8fcba 100644 --- a/alto/models/deepseek_v3/config_registry.py +++ b/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,7 +43,7 @@ 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 @@ -61,8 +61,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 diff --git a/alto/models/deepseek_v3/configs/lpt_recipe.yaml b/alto/models/deepseek_v3/configs/lpt_recipe.yaml index cd765eef..20794b74 100644 --- a/alto/models/deepseek_v3/configs/lpt_recipe.yaml +++ b/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 diff --git a/alto/models/gpt_oss/config_registry.py b/alto/models/gpt_oss/config_registry.py index fe51292b..c8ef15dd 100644 --- a/alto/models/gpt_oss/config_registry.py +++ b/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,7 +44,7 @@ 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 @@ -63,7 +63,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,7 +72,7 @@ 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 @@ -101,7 +101,7 @@ 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 diff --git a/alto/models/gpt_oss/configs/lpt_recipe.yaml b/alto/models/gpt_oss/configs/lpt_recipe.yaml index ad935b53..f2d78836 100644 --- a/alto/models/gpt_oss/configs/lpt_recipe.yaml +++ b/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 diff --git a/alto/models/llama3/config_registry.py b/alto/models/llama3/config_registry.py index 097d1256..83caa2de 100644 --- a/alto/models/llama3/config_registry.py +++ b/alto/models/llama3/config_registry.py @@ -54,12 +54,12 @@ def llama3_debugmodel() -> Trainer.Config: config = llama3_debugmodel_orig() - config.profiling.enable_profiling = False + config.profiler.enable_profiling = False config.training.steps = 0 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 @@ -85,13 +85,13 @@ def llama3_1b() -> Trainer.Config: config = llama3_1b_orig() config.hf_assets_path = "/group/archive_dataset_6_nobkup/archive_modelzoo/sequence_learning/weights/nlp-pretrained-model/meta-llama/Llama-3.2-1B" config.metrics.log_freq = 1 - config.profiling.enable_profiling = False + config.profiler.enable_profiling = False config.training.steps = 0 config.training.local_batch_size = 1 config.training.global_batch_size = 10 config.training.seq_len = 8192 config.dataloader.dataset = "c4_test" - config.activation_checkpoint.mode = "none" + config.activation_checkpoint = None config.checkpoint.enable = True config.checkpoint.interval = 10 config.checkpoint.initial_load_path = "/group/archive_dataset_6_nobkup/archive_modelzoo/sequence_learning/weights/nlp-pretrained-model/meta-llama/Llama-3.2-1B" @@ -128,7 +128,7 @@ def llama3_8b_pretrain() -> Trainer.Config: config.dump_folder = "llama3_8b-mi308-pretrain-subset-gbs384-lr1e-4-outputs" config.metrics.log_freq = 1 config.metrics.enable_tensorboard = True - config.profiling.enable_profiling = False + config.profiler.enable_profiling = False config.training.steps = 5000 config.training.local_batch_size = 2 config.training.global_batch_size = 384 @@ -143,7 +143,7 @@ def llama3_8b_pretrain() -> Trainer.Config: config.parallelism.expert_parallel_degree = 1 config.parallelism.expert_tensor_parallel_degree = 1 config.parallelism.tensor_parallel_degree = 1 - config.activation_checkpoint.mode = "none" + config.activation_checkpoint = None config.checkpoint.enable = False config.checkpoint.interval = 10 config.checkpoint.initial_load_path = "/huggingface/hub/models--unsloth--Llama-3.1-8B/snapshots/3f0d51f8e5640f98f1a96ea9044a0e55c0a83814" @@ -221,13 +221,13 @@ def llama3_8b() -> Trainer.Config: config = llama3_8b_orig() config.hf_assets_path = LLAMA3_8B_PATH config.metrics.log_freq = 1 - config.profiling.enable_profiling = False + config.profiler.enable_profiling = False config.training.steps = 0 config.training.local_batch_size = 1 config.training.global_batch_size = 8 config.training.seq_len = 2048 config.dataloader = HuggingFaceTextDataLoader.Config(dataset="c4_test") - config.activation_checkpoint.mode = "none" + config.activation_checkpoint = None config.checkpoint.enable = True config.checkpoint.interval = 10 config.checkpoint.initial_load_path = LLAMA3_8B_PATH @@ -385,13 +385,13 @@ def instella_3b() -> Trainer.Config: config = instella_3b_orig() config.hf_assets_path = "/group/ossmodelzoo/hanwang2/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e" config.metrics.log_freq = 1 - config.profiling.enable_profiling = False + config.profiler.enable_profiling = False config.training.steps = 0 config.training.local_batch_size = 1 config.training.global_batch_size = 10 config.training.seq_len = 4096 config.dataloader.dataset = "c4_test" - config.activation_checkpoint.mode = "none" + config.activation_checkpoint = None config.checkpoint.enable = True config.checkpoint.interval = 10 config.checkpoint.initial_load_path = "/group/ossmodelzoo/hanwang2/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e" diff --git a/alto/models/llama3/configs/admm_recipe.yaml b/alto/models/llama3/configs/admm_recipe.yaml index b9bdd6f8..416b7a92 100644 --- a/alto/models/llama3/configs/admm_recipe.yaml +++ b/alto/models/llama3/configs/admm_recipe.yaml @@ -11,4 +11,4 @@ sparsification_stage: admm_pruning_granularity: 32 admm_iterations: 32 targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/admm_structured_recipe.yaml b/alto/models/llama3/configs/admm_structured_recipe.yaml index 17501d64..a1a381ac 100644 --- a/alto/models/llama3/configs/admm_structured_recipe.yaml +++ b/alto/models/llama3/configs/admm_structured_recipe.yaml @@ -10,4 +10,4 @@ pruning_stage: admm_iterations: 16 dampening_frac: 0.01 targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/alps_recipe.yaml b/alto/models/llama3/configs/alps_recipe.yaml index c2530745..633b1600 100644 --- a/alto/models/llama3/configs/alps_recipe.yaml +++ b/alto/models/llama3/configs/alps_recipe.yaml @@ -12,4 +12,4 @@ sparsification_stage: alps_update_iter: 16 alps_switch_iter: 16 targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/awq_debug_recipe.yaml b/alto/models/llama3/configs/awq_debug_recipe.yaml index b329af00..4bc21c91 100644 --- a/alto/models/llama3/configs/awq_debug_recipe.yaml +++ b/alto/models/llama3/configs/awq_debug_recipe.yaml @@ -5,7 +5,7 @@ quantization_stage: AWQModifier: sequential_granularity: "sublayer" n_grid: 20 - ignore: ["output"] + ignore: ["lm_head"] mappings: - smooth_layer: "re:.*attention_norm" balance_layers: diff --git a/alto/models/llama3/configs/awq_recipe.yaml b/alto/models/llama3/configs/awq_recipe.yaml index d33c039a..0ab88397 100644 --- a/alto/models/llama3/configs/awq_recipe.yaml +++ b/alto/models/llama3/configs/awq_recipe.yaml @@ -6,7 +6,7 @@ quantization_stage: AWQModifier: sequential_granularity: "sublayer" n_grid: 20 - ignore: ["output"] + ignore: ["lm_head"] mappings: - smooth_layer: "re:.*attention_norm" balance_layers: diff --git a/alto/models/llama3/configs/cosine_similarity_recipe.yaml b/alto/models/llama3/configs/cosine_similarity_recipe.yaml index e9b4d399..827f7d3c 100644 --- a/alto/models/llama3/configs/cosine_similarity_recipe.yaml +++ b/alto/models/llama3/configs/cosine_similarity_recipe.yaml @@ -8,4 +8,4 @@ pruning_stage: pruning_dimension: "layer" targets: - "re:.*\\.wq$" - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/gptq_recipe.yaml b/alto/models/llama3/configs/gptq_recipe.yaml index 11721ead..5d162681 100644 --- a/alto/models/llama3/configs/gptq_recipe.yaml +++ b/alto/models/llama3/configs/gptq_recipe.yaml @@ -7,7 +7,7 @@ quantization_stage: sequential_granularity: "sublayer" block_size: 128 dampening_frac: 0.01 - ignore: ["output"] + ignore: ["lm_head"] config_groups: group_0: weights: diff --git a/alto/models/llama3/configs/guanchen_recipe.yaml b/alto/models/llama3/configs/guanchen_recipe.yaml index 6c25fc9b..f31232ae 100644 --- a/alto/models/llama3/configs/guanchen_recipe.yaml +++ b/alto/models/llama3/configs/guanchen_recipe.yaml @@ -4,7 +4,7 @@ # sparsity: 0.5 # mask_structure: "2:4" # targets: ["Linear"] -# ignore: ["output"] +# ignore: ["lm_head"] # sparsity_stage: # sparsity_modifiers: @@ -12,7 +12,7 @@ # sparsity: 0.5 # mask_structure: "2:4" # targets: ["Linear"] -# ignore: ["output"] +# ignore: ["lm_head"] # sparsity_stage: # sparsity_modifiers: @@ -20,7 +20,7 @@ # sparsity: 0.5 # mask_structure: "2:4" # targets: ["Linear"] -# ignore: ["output"] +# ignore: ["lm_head"] # sparsity_stage: # sparsity_modifiers: @@ -28,7 +28,7 @@ # sparsity: 0.5 # mask_structure: "2:4" # targets: ["Linear"] -# ignore: ["output"] +# ignore: ["lm_head"] pruning_stage: pruning_modifiers: diff --git a/alto/models/llama3/configs/lpt_recipe.yaml b/alto/models/llama3/configs/lpt_recipe.yaml index 2a12a1d1..3a6a36cd 100644 --- a/alto/models/llama3/configs/lpt_recipe.yaml +++ b/alto/models/llama3/configs/lpt_recipe.yaml @@ -3,7 +3,7 @@ training_stage: LowPrecisionTrainingModifier: scheme: "mxfp4" targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] use_2dblock_x: false use_2dblock_w: true use_hadamard: true diff --git a/alto/models/llama3/configs/magnitude_recipe.yaml b/alto/models/llama3/configs/magnitude_recipe.yaml index c9d4bec3..ed3ccef7 100644 --- a/alto/models/llama3/configs/magnitude_recipe.yaml +++ b/alto/models/llama3/configs/magnitude_recipe.yaml @@ -7,4 +7,4 @@ sparsification_stage: sparsity: 0.5 mask_structure: "0:0" targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/mx6_wa_recipe.yaml b/alto/models/llama3/configs/mx6_wa_recipe.yaml index 5ab4d3b1..d04d6d13 100644 --- a/alto/models/llama3/configs/mx6_wa_recipe.yaml +++ b/alto/models/llama3/configs/mx6_wa_recipe.yaml @@ -4,7 +4,7 @@ quantization_stage: quantization_modifiers: QuantizationModifier: - ignore: ["output"] # lm_head not quantized + ignore: ["lm_head"] # lm_head not quantized sequential: false config_groups: group_0: diff --git a/alto/models/llama3/configs/mx9_wa_recipe.yaml b/alto/models/llama3/configs/mx9_wa_recipe.yaml index 3df48414..93ea42b9 100644 --- a/alto/models/llama3/configs/mx9_wa_recipe.yaml +++ b/alto/models/llama3/configs/mx9_wa_recipe.yaml @@ -4,7 +4,7 @@ quantization_stage: quantization_modifiers: QuantizationModifier: - ignore: ["output"] # lm_head not quantized + ignore: ["lm_head"] # lm_head not quantized sequential: false config_groups: group_0: diff --git a/alto/models/llama3/configs/obs_recipe.yaml b/alto/models/llama3/configs/obs_recipe.yaml index 2447c806..0b214a9b 100644 --- a/alto/models/llama3/configs/obs_recipe.yaml +++ b/alto/models/llama3/configs/obs_recipe.yaml @@ -10,4 +10,4 @@ pruning_stage: attention_pruning_granularity: 1 dampening_frac: 0.01 targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/recipe.yaml b/alto/models/llama3/configs/recipe.yaml index bc4da067..50745533 100644 --- a/alto/models/llama3/configs/recipe.yaml +++ b/alto/models/llama3/configs/recipe.yaml @@ -4,11 +4,11 @@ # sparsity: 0.5 # mask_structure: "2:4" # targets: ["Linear"] -# ignore: ["output"] +# ignore: ["lm_head"] quantization_stage: quantization_modifiers: QuantizationModifier: - ignore: ["output"] + ignore: ["lm_head"] config_groups: group_0: weights: diff --git a/alto/models/llama3/configs/rtn_recipe.yaml b/alto/models/llama3/configs/rtn_recipe.yaml index 072ec777..8bfc334e 100644 --- a/alto/models/llama3/configs/rtn_recipe.yaml +++ b/alto/models/llama3/configs/rtn_recipe.yaml @@ -1,7 +1,7 @@ quantization_stage: quantization_modifiers: QuantizationModifier: - ignore: ["output"] + ignore: ["lm_head"] config_groups: group_0: weights: diff --git a/alto/models/llama3/configs/sparsegpt_recipe.yaml b/alto/models/llama3/configs/sparsegpt_recipe.yaml index f785b91c..e2785804 100644 --- a/alto/models/llama3/configs/sparsegpt_recipe.yaml +++ b/alto/models/llama3/configs/sparsegpt_recipe.yaml @@ -9,4 +9,4 @@ sparsification_stage: block_size: 128 dampening_frac: 0.01 targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/wanda_recipe.yaml b/alto/models/llama3/configs/wanda_recipe.yaml index 0f3aa082..de94fe02 100644 --- a/alto/models/llama3/configs/wanda_recipe.yaml +++ b/alto/models/llama3/configs/wanda_recipe.yaml @@ -7,4 +7,4 @@ sparsification_stage: sparsity: 0.5 mask_structure: "0:0" targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/llama3/configs/wanda_structured_recipe.yaml b/alto/models/llama3/configs/wanda_structured_recipe.yaml index b5541d9c..57791005 100644 --- a/alto/models/llama3/configs/wanda_structured_recipe.yaml +++ b/alto/models/llama3/configs/wanda_structured_recipe.yaml @@ -7,4 +7,4 @@ pruning_stage: sparsity: 0.25 pruning_dimension: "mlp" targets: ["Linear"] - ignore: ["output"] + ignore: ["lm_head"] diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 8c751354..7db81ba8 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -21,7 +21,7 @@ def patch(cls): cls._patched = True cls.patch_fake_quantize() - cls.patch_apply_rotary_emb_complex() + # cls.patch_apply_rotary_emb_complex() for model_name in SUPPORTED_MODELS: model_module = importlib.import_module(f"torchtitan.models.{model_name}") @@ -99,36 +99,36 @@ def fake_quantize(x, scale, zero_point, args, g_idx, global_scale): forward_module.fake_quantize = fake_quantize - @classmethod - def patch_apply_rotary_emb_complex(cls): - from torchtitan.models.common import rope, attention - original_apply_rotary_emb_complex = rope.apply_rotary_emb_complex - - def apply_rotary_emb_complex( - xq: torch.Tensor, - xk: torch.Tensor, - freqs_cis: torch.Tensor, - positions: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - head_dim = xq.shape[-1] - xq = xq.reshape( - *xq.shape[:-1], - 2, - head_dim // 2, - ).transpose(-1, -2).reshape( - *xq.shape[:-1], - head_dim, - ).contiguous() - xk = xk.reshape( - *xk.shape[:-1], - 2, - head_dim // 2, - ).transpose(-1, -2).reshape( - *xk.shape[:-1], - head_dim, - ).contiguous() - return original_apply_rotary_emb_complex(xq, xk, freqs_cis, positions) - - rope.apply_rotary_emb_complex = apply_rotary_emb_complex - attention.apply_rotary_emb_complex = apply_rotary_emb_complex - patch("torchtitan.models.common.rope.apply_rotary_emb_complex", apply_rotary_emb_complex).__enter__() + # @classmethod + # def patch_apply_rotary_emb_complex(cls): + # from torchtitan.models.common import rope, attention + # original_apply_rotary_emb_complex = rope.apply_rotary_emb_complex + + # def apply_rotary_emb_complex( + # xq: torch.Tensor, + # xk: torch.Tensor, + # freqs_cis: torch.Tensor, + # positions: torch.Tensor | None = None, + # ) -> tuple[torch.Tensor, torch.Tensor]: + # head_dim = xq.shape[-1] + # xq = xq.reshape( + # *xq.shape[:-1], + # 2, + # head_dim // 2, + # ).transpose(-1, -2).reshape( + # *xq.shape[:-1], + # head_dim, + # ).contiguous() + # xk = xk.reshape( + # *xk.shape[:-1], + # 2, + # head_dim // 2, + # ).transpose(-1, -2).reshape( + # *xk.shape[:-1], + # head_dim, + # ).contiguous() + # return original_apply_rotary_emb_complex(xq, xk, freqs_cis, positions) + + # rope.apply_rotary_emb_complex = apply_rotary_emb_complex + # attention.apply_rotary_emb_complex = apply_rotary_emb_complex + # patch("torchtitan.models.common.rope.apply_rotary_emb_complex", apply_rotary_emb_complex).__enter__() diff --git a/alto/modifiers/base.py b/alto/modifiers/base.py index b6b71697..33c190a1 100644 --- a/alto/modifiers/base.py +++ b/alto/modifiers/base.py @@ -9,12 +9,15 @@ # modified from https://github.com/vllm-project/llm-compressor/blob/f3f14af3ee56e35db7e1faf6da8833f84a570baf/src/llmcompressor/modifiers/modifier.py +from typing import TYPE_CHECKING from abc import abstractmethod from pydantic import ConfigDict import torch from torch.nn import Module from torchtitan.tools.logging import logger from alto.modifiers.utils import HooksMixin +if TYPE_CHECKING: + from torchtitan.protocols.model import BaseModel __all__ = ["Modifier"] @@ -52,6 +55,9 @@ def requires_training_mode(self) -> bool: def convert(self, model: Module, **kwargs) -> bool: return self.on_convert(model, **kwargs) + + def convert_config(self, model_config: "BaseModel.Config") -> bool: + return self.on_convert_config(model_config) def initialize(self, model_parts: list[Module], **kwargs): if self.initialized_: @@ -102,6 +108,10 @@ def on_post_step(self, model_parts: list[Module], **kwargs) -> bool: def on_convert(self, model: Module, **kwargs) -> bool: raise NotImplementedError + @abstractmethod + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: + raise NotImplementedError + def _infer_sequential_targets(self, model: torch.nn.Module) -> str | list[str]: match self.sequential_targets: case None: diff --git a/alto/modifiers/distillation/base.py b/alto/modifiers/distillation/base.py index 803cec01..f4e4fa10 100644 --- a/alto/modifiers/distillation/base.py +++ b/alto/modifiers/distillation/base.py @@ -4,7 +4,7 @@ from abc import abstractmethod from functools import partial -from typing import Any, Literal +from typing import Any, Literal, TYPE_CHECKING import gc import torch @@ -23,6 +23,8 @@ from alto.utils.pytorch.module import ( get_layers,) from alto.modifiers.distillation.utils import losses +if TYPE_CHECKING: + from torchtitan.protocols.model import BaseModel TEACHER_OBSERVER_BASE_NAME = "output" STUDENT_OBSERVER_BASE_NAME = "student_output" @@ -284,3 +286,6 @@ def _build_optimizers(self, model_parts: list[Module]): def on_convert(self, model: Module, **kwargs) -> bool: return True + + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: + return True diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 319d0835..64e7e46d 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -2,25 +2,27 @@ # # SPDX-License-Identifier: MIT -from typing import Literal +from typing import Literal, TYPE_CHECKING import torch from torch.nn import Module from compressed_tensors.utils import match_named_modules from pydantic import PrivateAttr, Field, field_validator, model_validator -from torchtitan.models.common.attention import BaseAttention -from torchtitan.models.common.moe.utils import set_token_group_alignment_size_m +from torchtitan.models.common.attention import BaseAttention, ScaledDotProductAttention from torchtitan.tools.logging import logger from alto.modifiers import Modifier from alto.kernels.dispatch import ( swap_params, TrainingOpConfig, - LPScaledDotProductAttentionWrapper, + LPScaledDotProductAttention, ) from alto.kernels.fp4.mxfp4.mxfp_grouped_gemm.autotune import ALIGN_SIZE_M from alto.nn import DecomposedLinear from alto.components.optimizer import DeOscillationConfig, enable_de_oscillation +if TYPE_CHECKING: + from torchtitan.protocols.model import BaseModel + __all__ = ["LowPrecisionTrainingModifier"] @@ -71,8 +73,6 @@ class LowPrecisionTrainingModifier(Modifier): def __init__(self, **kwargs): super().__init__(**kwargs) - set_token_group_alignment_size_m(ALIGN_SIZE_M) - @field_validator("targets", mode="before") def validate_targets(cls, value: str | list[str]) -> list[str]: if isinstance(value, str): @@ -166,8 +166,8 @@ 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): - assert module.attn_backend == "sdpa", "Only SDPA attention is supported for now." - module.inner_attention = LPScaledDotProductAttentionWrapper(config=scheme_obj) + assert isinstance(module.inner_attention, ScaledDotProductAttention), "Only SDPA attention is supported for now." + module.inner_attention = LPScaledDotProductAttention(config=scheme_obj) elif isinstance(module, torch.nn.Linear): if self.lora_rank > 0: module = DecomposedLinear.from_linear(module, lora_rank=self.lora_rank) @@ -184,6 +184,16 @@ def on_convert(self, model: Module, **kwargs) -> bool: logger.info(f"LowPrecisionTrainingModifier converted model: {model}") return True + + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: + # convert configs to enable token alignment + from torchtitan.models.common.moe import GroupedExperts + from torchtitan.components.quantization.utils import swap_token_dispatcher + + for _fqn, config, parent, attr in model_config.traverse(GroupedExperts.Config): + swap_token_dispatcher(config, ALIGN_SIZE_M) + raise NotImplementedError(f"GroupedExperts {_fqn}: {config}") + return True def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: for model_part in model_parts: diff --git a/alto/modifiers/pruning/base.py b/alto/modifiers/pruning/base.py index 9f229e04..fb965bb3 100644 --- a/alto/modifiers/pruning/base.py +++ b/alto/modifiers/pruning/base.py @@ -11,7 +11,7 @@ from abc import abstractmethod from functools import partial -from typing import Any, Generator, Literal +from typing import Any, Generator, Literal, TYPE_CHECKING import gc import re @@ -30,6 +30,8 @@ get_prunable_layers, match_targets, ) +if TYPE_CHECKING: + from torchtitan.protocols.model import BaseModel LAYER_OBSERVER_BASE_NAME = "pruning" PruningDim = Literal["attn", "mlp", "attn+mlp", "mlp+attn", "layer", "sublayer", "hidden_dim"] @@ -233,6 +235,9 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool: def on_convert(self, model: Module, **kwargs) -> bool: return True + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: + return True + # ---- sequential helpers ------------------------------------------- def _capture_hook(self, module: Module, args: Any, kwargs: Any = None): diff --git a/alto/modifiers/quantization/base.py b/alto/modifiers/quantization/base.py index 3511fdd2..a518a1f1 100644 --- a/alto/modifiers/quantization/base.py +++ b/alto/modifiers/quantization/base.py @@ -11,7 +11,7 @@ # Subclasses override _process_block() for algorithm-specific per-block logic. import gc -from typing import Any, Literal, Optional, Union +from typing import Any, Literal, Optional, Union, TYPE_CHECKING import torch import tqdm @@ -29,6 +29,9 @@ from alto.modifiers.quantization.calibration import update_weight_zp_scale from alto.utils.pytorch.module import get_layers +if TYPE_CHECKING: + from torchtitan.protocols.model import BaseModel + __all__ = ["QuantizationModifier"] @@ -112,6 +115,9 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool: def on_convert(self, model: Module, **kwargs) -> bool: return True + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: + return True + # ---- sequential loop (template) ----------------------------------- def _sequential_post_step(self, model_parts: list[Module]) -> bool: diff --git a/alto/modifiers/sparsification/base.py b/alto/modifiers/sparsification/base.py index 2e3d94ff..6d667b1a 100644 --- a/alto/modifiers/sparsification/base.py +++ b/alto/modifiers/sparsification/base.py @@ -11,7 +11,7 @@ from abc import abstractmethod from functools import partial -from typing import Any, Generator, Literal +from typing import Any, Generator, Literal, TYPE_CHECKING import gc import re @@ -30,6 +30,8 @@ get_prunable_layers, match_targets, ) +if TYPE_CHECKING: + from torchtitan.protocols.model import BaseModel LAYER_OBSERVER_BASE_NAME = "sparsity" @@ -242,6 +244,9 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool: def on_convert(self, model: Module, **kwargs) -> bool: return True + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: + return True + # ---- sequential helpers ------------------------------------------- def _capture_hook(self, module: Module, args: Any, kwargs: Any = None): diff --git a/alto/train.py b/alto/train.py index cec58aec..92978df3 100644 --- a/alto/train.py +++ b/alto/train.py @@ -148,9 +148,10 @@ def forward_step( model_parts = self.model_parts parallel_dims = self.parallel_dims - inputs, _, extra_inputs, extra_kwargs = self.post_dataloading_process(input_dict, labels) + inputs, labels, extra_kwargs = self.post_dataloading_process(input_dict, labels) if parallel_dims.pp_enabled: + loss_kwargs = {"global_valid_tokens": global_valid_tokens} targets, losses = None, None result = None with self.train_context(): @@ -158,16 +159,18 @@ def forward_step( if self.pp_has_first_stage: self.pp_schedule.eval( inputs, - **extra_inputs, **extra_kwargs, target=targets, losses=losses, + loss_kwargs=loss_kwargs, + return_outputs=False, ) elif self.pp_has_last_stage: result = self.pp_schedule.eval( **extra_kwargs, target=targets, losses=losses, + loss_kwargs=loss_kwargs, return_outputs=True, ) else: @@ -175,13 +178,15 @@ def forward_step( **extra_kwargs, target=targets, losses=losses, + loss_kwargs=loss_kwargs, + return_outputs=False, ) else: # Non-PP forward / backward with self.train_context(): assert len(model_parts) == 1 with self.maybe_enable_amp: - result = model_parts[0](inputs, **extra_inputs, **extra_kwargs) + result = model_parts[0](inputs, **extra_kwargs) return result diff --git a/alto/utils/exportation/export.py b/alto/utils/exportation/export.py index e728e47a..e803e4b9 100644 --- a/alto/utils/exportation/export.py +++ b/alto/utils/exportation/export.py @@ -60,17 +60,6 @@ def patched_finfo(dtype): torch.finfo = orig_finfo_func -def hot_fix_for_tied_word_embeddings(model: torch.nn.Module, compressor: ModelCompressor): - if not getattr(model.config, "enable_weight_tying", False): - return - # in the current impl, lm_head is not a linear layer, - # so we need to add it to the ignore list manually - if compressor.quantization_config is not None: - compressor.quantization_config.ignore.append("lm_head") - if compressor.sparsity_config is not None: - compressor.sparsity_config.ignore.append("lm_head") - - def get_model_compressor( model: torch.nn.Module, state_dict: dict[str, torch.Tensor], @@ -205,7 +194,6 @@ def convert_to_hf( compressor.quantization_config.ignore) if compressor.sparsity_config is not None: compressor.sparsity_config.ignore = sd_adapter.map_ignore_list_to_hf(compressor.sparsity_config.ignore) - hot_fix_for_tied_word_embeddings(model, compressor) state_dict = compressor.compress(model) compressor.update_config(output_dir) diff --git a/examples/ADMM_Llama_3.1_8B.sh b/examples/ADMM_Llama_3.1_8B.sh index ea64fee4..61c6f992 100755 --- a/examples/ADMM_Llama_3.1_8B.sh +++ b/examples/ADMM_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/ALPS_Llama_3.1_8B.sh b/examples/ALPS_Llama_3.1_8B.sh index 7e839754..208f83df 100755 --- a/examples/ALPS_Llama_3.1_8B.sh +++ b/examples/ALPS_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/AWQ_Llama_3.1_8B.sh b/examples/AWQ_Llama_3.1_8B.sh index dde18ae4..4fcf6ebd 100755 --- a/examples/AWQ_Llama_3.1_8B.sh +++ b/examples/AWQ_Llama_3.1_8B.sh @@ -29,6 +29,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/AdmmStructured_Llama_3.1_8B.sh b/examples/AdmmStructured_Llama_3.1_8B.sh index 472d28d9..54ab2992 100755 --- a/examples/AdmmStructured_Llama_3.1_8B.sh +++ b/examples/AdmmStructured_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/CosineSimilarity_Llama_3.1_8B.sh b/examples/CosineSimilarity_Llama_3.1_8B.sh index 0dc6e180..335f4a4d 100755 --- a/examples/CosineSimilarity_Llama_3.1_8B.sh +++ b/examples/CosineSimilarity_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/GPTQ_Llama_3.1_8B.sh b/examples/GPTQ_Llama_3.1_8B.sh index 1296fd9f..22733aa4 100755 --- a/examples/GPTQ_Llama_3.1_8B.sh +++ b/examples/GPTQ_Llama_3.1_8B.sh @@ -28,6 +28,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/Magnitude_Llama_3.1_8B.sh b/examples/Magnitude_Llama_3.1_8B.sh index d6e37870..b00bcae4 100755 --- a/examples/Magnitude_Llama_3.1_8B.sh +++ b/examples/Magnitude_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/OBS_Llama_3.1_8B.sh b/examples/OBS_Llama_3.1_8B.sh index 29bc5bbd..93c64eee 100755 --- a/examples/OBS_Llama_3.1_8B.sh +++ b/examples/OBS_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/RTN_Llama_3.1_8B.sh b/examples/RTN_Llama_3.1_8B.sh index a9be4df7..7dbe1c5b 100755 --- a/examples/RTN_Llama_3.1_8B.sh +++ b/examples/RTN_Llama_3.1_8B.sh @@ -41,6 +41,7 @@ if [ -n "$COMM_MODE" ]; then else # Normal training with torchrun PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/SparseGPT_Llama_3.1_8B.sh b/examples/SparseGPT_Llama_3.1_8B.sh index 075c48db..2f708cb5 100755 --- a/examples/SparseGPT_Llama_3.1_8B.sh +++ b/examples/SparseGPT_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/WandaStructured_Llama_3.1_8B.sh b/examples/WandaStructured_Llama_3.1_8B.sh index 385740fb..4d47fae3 100755 --- a/examples/WandaStructured_Llama_3.1_8B.sh +++ b/examples/WandaStructured_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/Wanda_Llama_3.1_8B.sh b/examples/Wanda_Llama_3.1_8B.sh index 4547a654..c7fa9301 100755 --- a/examples/Wanda_Llama_3.1_8B.sh +++ b/examples/Wanda_Llama_3.1_8B.sh @@ -26,6 +26,7 @@ if [ -n "$COMM_MODE" ]; then NGPU="${NGPU}" LOCAL_RANK=0 python3 -m ${TRAIN_FILE} --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1 else PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ diff --git a/examples/llama3.2_1b_mx9.sh b/examples/llama3.2_1b_mx9.sh index 204794a6..59bb221f 100755 --- a/examples/llama3.2_1b_mx9.sh +++ b/examples/llama3.2_1b_mx9.sh @@ -49,6 +49,7 @@ if [ -n "$COMM_MODE" ]; then else PYTORCH_ALLOC_CONF="expandable_segments:True" \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ + HSA_NO_SCRATCH_RECLAIM=1 \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ -m ${TRAIN_FILE} \ diff --git a/examples/run.sh b/examples/run.sh index 4c6ff0c6..2957492a 100755 --- a/examples/run.sh +++ b/examples/run.sh @@ -44,6 +44,7 @@ if [ -n "$COMM_MODE" ]; then else # Normal training with torchrun PYTORCH_ALLOC_CONF="expandable_segments:True" \ + HSA_NO_SCRATCH_RECLAIM=1 \ TORCHFT_LIGHTHOUSE=${TORCHFT_LIGHTHOUSE} \ torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \ --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \ From 3a46ed350004491f3b8a5b9063458fd259af8cff Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 14 Jul 2026 07:40:59 +0000 Subject: [PATCH 02/14] fix: apply_rotary_emb_complex patch --- 3rdparty/torchtitan | 2 +- alto/models/llama3/config_registry.py | 36 +++++++-------- alto/models/patcher.py | 66 +++++++++++++-------------- alto/modifiers/distillation/base.py | 20 +++++--- alto/train.py | 3 +- tests/integration/instella_3b_opt.sh | 11 ----- version.txt | 2 +- 7 files changed, 66 insertions(+), 74 deletions(-) delete mode 100755 tests/integration/instella_3b_opt.sh diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan index 8296b499..b49df50f 160000 --- a/3rdparty/torchtitan +++ b/3rdparty/torchtitan @@ -1 +1 @@ -Subproject commit 8296b49912ddd98b75a0584e15ce04ff0fc9a9db +Subproject commit b49df50fc88ee9cc7ae68ba854063226b9697b4d diff --git a/alto/models/llama3/config_registry.py b/alto/models/llama3/config_registry.py index 83caa2de..5197b60e 100644 --- a/alto/models/llama3/config_registry.py +++ b/alto/models/llama3/config_registry.py @@ -2,7 +2,7 @@ # # SPDX-License-Identifier: MIT -from torchtitan.components.optimizer import OptimizersContainer +from torchtitan.components.optimizer import default_adamw from torchtitan.trainer import Trainer from torchtitan.protocols.model_converter import ModelConvertersContainer from torchtitan.models.llama3.config_registry import ( @@ -107,7 +107,7 @@ def llama3_1b() -> Trainer.Config: def llama3_1b_opt() -> Trainer.Config: config = llama3_1b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/recipe.yaml",), ],) @@ -177,7 +177,7 @@ def llama3_8b_lpt() -> Trainer.Config: def llama3_1b_gptq() -> Trainer.Config: config = llama3_1b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/gptq_recipe.yaml",), ],) @@ -187,7 +187,7 @@ def llama3_1b_gptq() -> Trainer.Config: def llama3_1b_awq() -> Trainer.Config: config = llama3_1b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/awq_recipe.yaml",), ],) @@ -197,7 +197,7 @@ def llama3_1b_awq() -> Trainer.Config: def llama3_1b_mx9_wa() -> Trainer.Config: config = llama3_1b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/mx9_wa_recipe.yaml",), ],) @@ -207,7 +207,7 @@ def llama3_1b_mx9_wa() -> Trainer.Config: def llama3_1b_mx6_wa() -> Trainer.Config: config = llama3_1b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/mx6_wa_recipe.yaml",), ],) @@ -243,7 +243,7 @@ def llama3_8b() -> Trainer.Config: def llama3_8b_gptq() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.training.global_batch_size = 128 config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/gptq_recipe.yaml",), @@ -254,7 +254,7 @@ def llama3_8b_gptq() -> Trainer.Config: def llama3_8b_rtn() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/rtn_recipe.yaml",), ],) @@ -264,7 +264,7 @@ def llama3_8b_rtn() -> Trainer.Config: def llama3_8b_awq() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/awq_recipe.yaml",), ],) @@ -279,7 +279,7 @@ def llama3_8b_awq() -> Trainer.Config: def llama3_8b_wanda() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/wanda_recipe.yaml",), ],) @@ -289,7 +289,7 @@ def llama3_8b_wanda() -> Trainer.Config: def llama3_8b_sparsegpt() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.training.global_batch_size = 128 config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/sparsegpt_recipe.yaml",), @@ -300,7 +300,7 @@ def llama3_8b_sparsegpt() -> Trainer.Config: def llama3_8b_magnitude() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/magnitude_recipe.yaml",), ],) @@ -310,7 +310,7 @@ def llama3_8b_magnitude() -> Trainer.Config: def llama3_8b_admm() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.training.global_batch_size = 128 config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/admm_recipe.yaml",), @@ -321,7 +321,7 @@ def llama3_8b_admm() -> Trainer.Config: def llama3_8b_alps() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.training.global_batch_size = 128 config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/alps_recipe.yaml",), @@ -337,7 +337,7 @@ def llama3_8b_alps() -> Trainer.Config: def llama3_8b_wanda_structured() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/wanda_structured_recipe.yaml",), ],) @@ -347,7 +347,7 @@ def llama3_8b_wanda_structured() -> Trainer.Config: def llama3_8b_obs() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.training.global_batch_size = 128 config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/obs_recipe.yaml",), @@ -358,7 +358,7 @@ def llama3_8b_obs() -> Trainer.Config: def llama3_8b_admm_structured() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.training.global_batch_size = 128 config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/admm_structured_recipe.yaml",), @@ -369,7 +369,7 @@ def llama3_8b_admm_structured() -> Trainer.Config: def llama3_8b_cosine_similarity() -> Trainer.Config: config = llama3_8b() config.training.steps = 1 - config.optimizer = OptimizersContainer.Config(lr=0.0) + config.optimizer = default_adamw(lr=0.0) config.model_converters = ModelConvertersContainer.Config(converters=[ ModelOptConverter.Config(recipe="./alto/models/llama3/configs/cosine_similarity_recipe.yaml",), ],) diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 7db81ba8..9497e543 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -21,7 +21,7 @@ def patch(cls): cls._patched = True cls.patch_fake_quantize() - # cls.patch_apply_rotary_emb_complex() + cls.patch_apply_rotary_emb_complex() for model_name in SUPPORTED_MODELS: model_module = importlib.import_module(f"torchtitan.models.{model_name}") @@ -99,36 +99,34 @@ def fake_quantize(x, scale, zero_point, args, g_idx, global_scale): forward_module.fake_quantize = fake_quantize - # @classmethod - # def patch_apply_rotary_emb_complex(cls): - # from torchtitan.models.common import rope, attention - # original_apply_rotary_emb_complex = rope.apply_rotary_emb_complex - - # def apply_rotary_emb_complex( - # xq: torch.Tensor, - # xk: torch.Tensor, - # freqs_cis: torch.Tensor, - # positions: torch.Tensor | None = None, - # ) -> tuple[torch.Tensor, torch.Tensor]: - # head_dim = xq.shape[-1] - # xq = xq.reshape( - # *xq.shape[:-1], - # 2, - # head_dim // 2, - # ).transpose(-1, -2).reshape( - # *xq.shape[:-1], - # head_dim, - # ).contiguous() - # xk = xk.reshape( - # *xk.shape[:-1], - # 2, - # head_dim // 2, - # ).transpose(-1, -2).reshape( - # *xk.shape[:-1], - # head_dim, - # ).contiguous() - # return original_apply_rotary_emb_complex(xq, xk, freqs_cis, positions) - - # rope.apply_rotary_emb_complex = apply_rotary_emb_complex - # attention.apply_rotary_emb_complex = apply_rotary_emb_complex - # patch("torchtitan.models.common.rope.apply_rotary_emb_complex", apply_rotary_emb_complex).__enter__() + @classmethod + def patch_apply_rotary_emb_complex(cls): + from torchtitan.models.common.rope import ComplexRoPE + original_apply_rotary_emb_complex = ComplexRoPE.apply_rotary_emb + + def apply_rotary_emb_complex( + cls, + xq: torch.Tensor, + xk: torch.Tensor, + rope_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + head_dim = xq.shape[-1] + xq = xq.reshape( + *xq.shape[:-1], + 2, + head_dim // 2, + ).transpose(-1, -2).reshape( + *xq.shape[:-1], + head_dim, + ).contiguous() + xk = xk.reshape( + *xk.shape[:-1], + 2, + head_dim // 2, + ).transpose(-1, -2).reshape( + *xk.shape[:-1], + head_dim, + ).contiguous() + return original_apply_rotary_emb_complex(xq, xk, rope_cache) + + ComplexRoPE.apply_rotary_emb = apply_rotary_emb_complex diff --git a/alto/modifiers/distillation/base.py b/alto/modifiers/distillation/base.py index f4e4fa10..ec7a7a5b 100644 --- a/alto/modifiers/distillation/base.py +++ b/alto/modifiers/distillation/base.py @@ -14,7 +14,7 @@ from torchtitan.tools.logging import logger from torchtitan.tools.utils import device_type from torchtitan.components.loss import IGNORE_INDEX -from torchtitan.components.optimizer import OptimizersContainer +from torchtitan.components.optimizer import OptimizersContainer, ParamGroupConfig from torchtitan.components.lr_scheduler import LRSchedulersContainer from alto.observers import Observer @@ -273,12 +273,18 @@ def _get_loss_fn(self, name: str): def _build_optimizers(self, model_parts: list[Module]): config = OptimizersContainer.Config( - name=self.optimizer, - lr=self.lr, - beta1=self.beta1, - beta2=self.beta2, - eps=self.eps, - weight_decay=self.weight_decay, + param_groups=[ + ParamGroupConfig( + pattern=r".*", + optimizer_name=self.optimizer, + optimizer_kwargs={ + "lr": self.lr, + "betas": (self.beta1, self.beta2), + "eps": self.eps, + "weight_decay": self.weight_decay, + }, + ) + ], implementation=self.implementation, ) diff --git a/alto/train.py b/alto/train.py index 92978df3..0b12a350 100644 --- a/alto/train.py +++ b/alto/train.py @@ -185,8 +185,7 @@ def forward_step( # Non-PP forward / backward with self.train_context(): assert len(model_parts) == 1 - with self.maybe_enable_amp: - result = model_parts[0](inputs, **extra_kwargs) + result = model_parts[0](inputs, **extra_kwargs) return result diff --git a/tests/integration/instella_3b_opt.sh b/tests/integration/instella_3b_opt.sh deleted file mode 100755 index c7d4ed13..00000000 --- a/tests/integration/instella_3b_opt.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash - -SCRIPT_DIR=$(dirname "$0") -cd $SCRIPT_DIR/../.. - -NGPU=1 \ -MODULE=llama3 \ -CONFIG=instella_3b_opt \ -./examples/run.sh \ - --hf_assets_path=/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e \ - --checkpoint.initial_load_path=/huggingface/hub/models--amd--Instella-3B-Stage1/snapshots/cb33253ab0a5b9f2ea0b98f3edd818d46454580e diff --git a/version.txt b/version.txt index 22701ea2..206c0852 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -0.0.2-dev0 +0.1.0-dev0 From b9786c95270bffef333dad5c48fbae9c2d8e87c7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 14 Jul 2026 08:45:39 +0000 Subject: [PATCH 03/14] fix: token dispatcher with paded tokens; fix: decomposed linear swap --- alto/components/converter.py | 8 +++ alto/modifiers/lpt/base.py | 43 +++++++++++++--- alto/nn/decomposed_linear.py | 55 ++++++++++----------- tests/integration/gpt_oss_debugmodel_lpt.sh | 3 +- 4 files changed, 72 insertions(+), 37 deletions(-) diff --git a/alto/components/converter.py b/alto/components/converter.py index 88ecb983..cd43b637 100644 --- a/alto/components/converter.py +++ b/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: diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 64e7e46d..849819a9 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -2,10 +2,10 @@ # # SPDX-License-Identifier: MIT -from typing import Literal, TYPE_CHECKING +from typing import Literal, Iterable, TYPE_CHECKING import torch from torch.nn import Module -from compressed_tensors.utils import match_named_modules +from compressed_tensors.utils import match_named_modules, match_name from pydantic import PrivateAttr, Field, field_validator, model_validator from torchtitan.models.common.attention import BaseAttention, ScaledDotProductAttention from torchtitan.tools.logging import logger @@ -184,22 +184,49 @@ def on_convert(self, model: Module, **kwargs) -> bool: logger.info(f"LowPrecisionTrainingModifier converted model: {model}") return True - + def on_convert_config(self, model_config: "BaseModel.Config") -> bool: # convert configs to enable token alignment from torchtitan.models.common.moe import GroupedExperts from torchtitan.components.quantization.utils import swap_token_dispatcher + from torchtitan.models.common.nn_modules import Linear for _fqn, config, parent, attr in model_config.traverse(GroupedExperts.Config): swap_token_dispatcher(config, ALIGN_SIZE_M) - raise NotImplementedError(f"GroupedExperts {_fqn}: {config}") + + def is_match( + name: str, + module_cls_name: str, + targets: str | Iterable[str], + ignore: str | Iterable[str] = tuple(), + ) -> bool: + targets = [targets] if isinstance(targets, str) else targets + ignore = [ignore] if isinstance(ignore, str) else ignore + + return any( + match_name(name, target) or module_cls_name == target + for target in targets + ) and not any( + match_name(name, ign) or module_cls_name == ign for ign in ignore + ) + + # replace Linear with DecomposedLinear + if self.lora_rank > 0: + resolved_targets = list(self.resolved_config.values()) + for _fqn, config, parent, attr in model_config.traverse(Linear.Config): + #raise ValueError(f"[{_fqn}] cfg={config},\nparent={parent},\nattr={attr}") + for target in resolved_targets: + if is_match(_fqn, "Linear", target, self.ignore): + 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) return True def on_initialize(self, model_parts: list[Module], **kwargs) -> bool: - for model_part in model_parts: - for child in model_part.modules(): - if isinstance(child, DecomposedLinear): - child.init_lora_weights(init_std=0.02) return True def on_pre_step(self, model_parts: list[Module], **kwargs) -> bool: diff --git a/alto/nn/decomposed_linear.py b/alto/nn/decomposed_linear.py index 54e721a2..3b2f12cb 100644 --- a/alto/nn/decomposed_linear.py +++ b/alto/nn/decomposed_linear.py @@ -1,23 +1,39 @@ # Copyright (c) 2026 Advanced Micro Devices, Inc. # # SPDX-License-Identifier: MIT - +from dataclasses import dataclass +from functools import partial import torch import torch.nn as nn import torch.nn.functional as F +from torchtitan.protocols.module import Module + +class DecomposedLinear(Module): + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + in_features: int + out_features: int + bias: bool = False + lora_rank: int = 32 + + _EXTRA_INIT = { + "u": partial(nn.init.normal_, std=0.02), + "v": nn.init.zeros_, + "sigma": nn.init.ones_, + } -class DecomposedLinear(nn.Module): - def __init__(self, in_features, out_features, bias=True, lora_rank=32): - super(DecomposedLinear, self).__init__() - self.in_features = in_features - self.out_features = out_features + def __init__(self, config: Config): + super().__init__() + self.in_features = config.in_features + self.out_features = config.out_features - self.weight = nn.Parameter(torch.empty(out_features, in_features)) - self.bias = nn.Parameter(torch.empty(out_features)) if bias else None - self.u = nn.Parameter(torch.empty(lora_rank, out_features)) # transposed - self.v = nn.Parameter(torch.empty(in_features, lora_rank)) - self.sigma = nn.Parameter(torch.empty(lora_rank)) + self.weight = nn.Parameter(torch.empty(self.out_features, self.in_features)) + self.bias = nn.Parameter(torch.empty(self.out_features)) if config.bias else None + self.u = nn.Parameter(torch.empty(config.lora_rank, self.out_features)) # transposed + self.v = nn.Parameter(torch.empty(self.in_features, config.lora_rank)) + self.sigma = nn.Parameter(torch.empty(config.lora_rank)) def forward(self, input): lora_update = (input @ self.v) * self.sigma @@ -25,20 +41,3 @@ def forward(self, input): if self.bias is not None: y += self.bias return y - - @classmethod - def from_linear(cls, linear: nn.Linear, lora_rank: int = 32): - new_layer = cls(linear.in_features, linear.out_features, linear.bias is not None, lora_rank) - new_layer.weight = linear.weight - new_layer.bias = linear.bias - device = linear.weight.device - dtype = linear.weight.dtype - new_layer.u.data = new_layer.u.data.to(device=device, dtype=dtype) - new_layer.v.data = new_layer.v.data.to(device=device, dtype=dtype) - new_layer.sigma.data = new_layer.sigma.data.to(device=device, dtype=dtype) - return new_layer - - def init_lora_weights(self, init_std: float = 0.02): - nn.init.normal_(self.u, mean=0.0, std=init_std) - nn.init.zeros_(self.v) - nn.init.ones_(self.sigma) diff --git a/tests/integration/gpt_oss_debugmodel_lpt.sh b/tests/integration/gpt_oss_debugmodel_lpt.sh index 7037a540..fb568804 100755 --- a/tests/integration/gpt_oss_debugmodel_lpt.sh +++ b/tests/integration/gpt_oss_debugmodel_lpt.sh @@ -6,4 +6,5 @@ cd $SCRIPT_DIR/../.. NGPU=2 \ MODULE=gpt_oss \ CONFIG=gpt_oss_debugmodel_lpt \ -./examples/run.sh +./examples/run.sh \ + --parallelism.expert_parallel_degree 2 From a9efbdbcf666f30f6e89bb2315aa63b61d289cc2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 15 Jul 2026 02:11:23 +0000 Subject: [PATCH 04/14] feat: basic DTensor support in wanda --- 3rdparty/torchtitan | 2 +- CHANGELOG.md | 7 +++++++ alto/modifiers/quantization/base.py | 21 ++++++++++++++++----- alto/modifiers/sparsification/wanda.py | 16 ++++++++++++++++ alto/observers/per_channel_norm.py | 17 ++++++++++++++--- 5 files changed, 54 insertions(+), 9 deletions(-) diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan index b49df50f..389944c5 160000 --- a/3rdparty/torchtitan +++ b/3rdparty/torchtitan @@ -1 +1 @@ -Subproject commit b49df50fc88ee9cc7ae68ba854063226b9697b4d +Subproject commit 389944c5cf665402a6d04fbbc58c7c4b77753ec9 diff --git a/CHANGELOG.md b/CHANGELOG.md index 4627965c..46a45e3e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/alto/modifiers/quantization/base.py b/alto/modifiers/quantization/base.py index a518a1f1..185f3217 100644 --- a/alto/modifiers/quantization/base.py +++ b/alto/modifiers/quantization/base.py @@ -73,11 +73,6 @@ class QuantizationModifier(Modifier, QuantizationMixin): # ---- lifecycle ---------------------------------------------------- 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 @@ -113,6 +108,22 @@ def on_finalize(self, model_parts: list[Module], **kwargs) -> bool: return True def on_convert(self, model: Module, **kwargs) -> bool: + if not QuantizationMixin.has_config(self): + raise ValueError("QuantizationModifier requires that quantization fields be specified") + # Note: qparams will be registered in the initialize_quantization method, + # so it has to be done before applying parallelism + QuantizationMixin.initialize_quantization(self, model) + + # patch param_init dict because the qparams are registered on meta device + # and need to be initialized after to_empty copy. + for mod_name, mod in model.named_modules(): + qscheme = getattr(mod, "quantization_scheme", None) + if qscheme is not None: + for pname, p in mod.named_parameters(): + if pname.endswith("scale"): + mod._param_init[pname] = torch.nn.init.ones_ + elif pname.endswith("zero_point"): + mod._param_init[pname] = torch.nn.init.zeros_ return True def on_convert_config(self, model_config: "BaseModel.Config") -> bool: diff --git a/alto/modifiers/sparsification/wanda.py b/alto/modifiers/sparsification/wanda.py index 24c62df2..10abf04c 100644 --- a/alto/modifiers/sparsification/wanda.py +++ b/alto/modifiers/sparsification/wanda.py @@ -11,6 +11,7 @@ import torch from torch.nn import Module +from torch.distributed.tensor import DTensor, Replicate from torchtitan.tools.logging import logger from alto.modifiers.sparsification.base import SparsityModifierBase @@ -79,6 +80,11 @@ def compress_modules(self): prune_n=self._prune_n, prune_m=self._prune_m, ) + if isinstance(module.weight, DTensor): + sparsified_weight = sparsified_weight.redistribute( + module.weight.device_mesh, + placements=module.weight.placements, + ) module.weight.data.copy_(sparsified_weight) def _sparsify_weight( @@ -111,6 +117,12 @@ def _sparsify_weight( W = W.to(dtype=PRECISION) S = row_scalar + + device_mesh = None + if isinstance(W, DTensor): + device_mesh = W.device_mesh + W = W.redistribute(device_mesh, placements=(Replicate(),)).to_local() + S = S.redistribute(device_mesh, placements=(Replicate(),)).to_local() W_metric = torch.abs(W) * torch.sqrt(S.reshape((1, -1))) @@ -138,4 +150,8 @@ def _sparsify_weight( W = W.reshape(final_shape).to(final_dtype) + if device_mesh is not None: + W = DTensor.from_local(W, device_mesh=device_mesh, placements=(Replicate(),)) + W_mask = DTensor.from_local(W_mask, device_mesh=device_mesh, placements=(Replicate(),)) + return W, W_mask.reshape(final_shape) diff --git a/alto/observers/per_channel_norm.py b/alto/observers/per_channel_norm.py index 6a1cda79..7c49c126 100644 --- a/alto/observers/per_channel_norm.py +++ b/alto/observers/per_channel_norm.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: MIT import torch +from torch.distributed.tensor import DTensor, Partial from alto.utils.pytorch.module import TransformerConv1D from .base import Observer, register_observer @@ -25,7 +26,10 @@ def __init__(self, *args, **kwargs) -> None: def make_empty_row_scalars(self) -> torch.Tensor: weight = self.module().weight num_columns = weight.shape[1] - return torch.zeros(num_columns, device=self.device) + S = torch.zeros(num_columns, device=self.device) + if isinstance(weight, DTensor): + S = DTensor.from_local(S, device_mesh=weight.device_mesh, placements=(Partial("avg"),)) + return S def get_current_min_max(self, observed: torch.Tensor): pass @@ -38,6 +42,13 @@ def forward_inner(self, x_orig): return x_orig with torch.no_grad(): + S = self.stats + if isinstance(S, DTensor): + S = S.to_local() + + # TODO: support TP + assert not isinstance(x_orig, DTensor), "TP is not supported for per_channel_norm observer" + module = self.module() inp = x_orig.detach() if inp.dim() == 2: @@ -62,11 +73,11 @@ def forward_inner(self, x_orig): inp = inp.permute([1, 0, 2]) inp = inp.flatten(1) - self.stats *= self.num_samples / (self.num_samples + num_added) + S *= self.num_samples / (self.num_samples + num_added) self.num_samples += num_added inp = inp.type(PRECISION) - self.stats += torch.norm(inp, p=2, dim=1)**2 / self.num_samples + S += torch.norm(inp, p=2, dim=1)**2 / self.num_samples return x_orig From ff64bebb2d6b105b182d6f24179520d0f7a10e51 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 15 Jul 2026 02:37:08 +0000 Subject: [PATCH 05/14] fix: restore from_linear() classmethod and update test constructor usage --- alto/nn/decomposed_linear.py | 14 ++++++++++++++ tests/unittest/nn/test_decomposed_linear.py | 2 +- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/alto/nn/decomposed_linear.py b/alto/nn/decomposed_linear.py index 3b2f12cb..332677ae 100644 --- a/alto/nn/decomposed_linear.py +++ b/alto/nn/decomposed_linear.py @@ -35,6 +35,20 @@ def __init__(self, config: Config): self.v = nn.Parameter(torch.empty(self.in_features, config.lora_rank)) self.sigma = nn.Parameter(torch.empty(config.lora_rank)) + @classmethod + def from_linear(cls, linear: nn.Linear, lora_rank: int = 32) -> "DecomposedLinear": + config = cls.Config( + in_features=linear.in_features, + out_features=linear.out_features, + bias=linear.bias is not None, + lora_rank=lora_rank, + ) + module = cls(config) + module.weight.data.copy_(linear.weight.data) + if linear.bias is not None: + module.bias.data.copy_(linear.bias.data) + return module + def forward(self, input): lora_update = (input @ self.v) * self.sigma y = F.linear(input, self.weight) + lora_update @ self.u diff --git a/tests/unittest/nn/test_decomposed_linear.py b/tests/unittest/nn/test_decomposed_linear.py index 99376d85..191540fa 100644 --- a/tests/unittest/nn/test_decomposed_linear.py +++ b/tests/unittest/nn/test_decomposed_linear.py @@ -41,7 +41,7 @@ def test_decomposed_linear(in_features, out_features, bias, lora_rank): @pytest.mark.parametrize("precision", ["mxfp4"]) def test_decomposed_linear_quantization(in_features, out_features, bias, lora_rank, precision): STD = 0.1 - decomposed_linear = DecomposedLinear(in_features, out_features, bias, lora_rank).to("cuda") + decomposed_linear = DecomposedLinear(DecomposedLinear.Config(in_features=in_features, out_features=out_features, bias=bias, lora_rank=lora_rank)).to("cuda") decomposed_linear.weight.data.normal_(mean=0, std=STD) if bias: decomposed_linear.bias.data.normal_(mean=0, std=STD) From 06e4f2fdee0e44a83e855e22b1605fb190c74e0d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 17 Jul 2026 01:46:13 -0500 Subject: [PATCH 06/14] fix: gpt_oss_20b optimizer and dataloader issues --- 3rdparty/torchtitan | 2 +- alto/models/deepseek_v3/config_registry.py | 1 - alto/models/gpt_oss/config_registry.py | 12 ++++-------- alto/models/llama3/config_registry.py | 3 +-- 4 files changed, 6 insertions(+), 12 deletions(-) diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan index 389944c5..cfa631dd 160000 --- a/3rdparty/torchtitan +++ b/3rdparty/torchtitan @@ -1 +1 @@ -Subproject commit 389944c5cf665402a6d04fbbc58c7c4b77753ec9 +Subproject commit cfa631dd08919107867a65694da32b3faf4fd46d diff --git a/alto/models/deepseek_v3/config_registry.py b/alto/models/deepseek_v3/config_registry.py index 14e8fcba..a1c7c90d 100644 --- a/alto/models/deepseek_v3/config_registry.py +++ b/alto/models/deepseek_v3/config_registry.py @@ -51,7 +51,6 @@ def deepseek_v3_16b() -> Trainer.Config: 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" diff --git a/alto/models/gpt_oss/config_registry.py b/alto/models/gpt_oss/config_registry.py index c8ef15dd..92232d4f 100644 --- a/alto/models/gpt_oss/config_registry.py +++ b/alto/models/gpt_oss/config_registry.py @@ -52,7 +52,6 @@ def gpt_oss_20b() -> Trainer.Config: 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/" @@ -77,11 +76,9 @@ def gpt_oss_20b_pretrain() -> Trainer.Config: 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 @@ -108,7 +104,7 @@ def gpt_oss_20b_pretrain() -> Trainer.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",), ],) diff --git a/alto/models/llama3/config_registry.py b/alto/models/llama3/config_registry.py index 5197b60e..1d9acf2a 100644 --- a/alto/models/llama3/config_registry.py +++ b/alto/models/llama3/config_registry.py @@ -133,7 +133,7 @@ def llama3_8b_pretrain() -> Trainer.Config: config.training.local_batch_size = 2 config.training.global_batch_size = 384 config.training.seq_len = 8192 - config.optimizer.lr = 1e-4 + config.optimizer.param_groups[0].optimizer_kwargs["lr"] = 1e-4 config.lr_scheduler.min_lr_factor = 0.0 config.lr_scheduler.warmup_steps = 500 config.lr_scheduler.decay_ratio = 0.9 @@ -141,7 +141,6 @@ def llama3_8b_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 = 1 - config.parallelism.expert_tensor_parallel_degree = 1 config.parallelism.tensor_parallel_degree = 1 config.activation_checkpoint = None config.checkpoint.enable = False From e84bc2d6e3d7d1bf522e28be85546d85df7aac21 Mon Sep 17 00:00:00 2001 From: Hann Wang Date: Fri, 17 Jul 2026 14:54:24 +0800 Subject: [PATCH 07/14] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- alto/nn/decomposed_linear.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/alto/nn/decomposed_linear.py b/alto/nn/decomposed_linear.py index 332677ae..c5e5bfaf 100644 --- a/alto/nn/decomposed_linear.py +++ b/alto/nn/decomposed_linear.py @@ -43,10 +43,23 @@ def from_linear(cls, linear: nn.Linear, lora_rank: int = 32) -> "DecomposedLinea bias=linear.bias is not None, lora_rank=lora_rank, ) - module = cls(config) - module.weight.data.copy_(linear.weight.data) - if linear.bias is not None: - module.bias.data.copy_(linear.bias.data) + module = cls(config).to(device=linear.weight.device, dtype=linear.weight.dtype) + + # Meta tensors cannot be copied-from; keep original params and rely on later init. + if linear.weight.is_meta: + module.weight = linear.weight + module.bias = linear.bias + else: + module.weight.data.copy_(linear.weight.data) + if linear.bias is not None: + module.bias.data.copy_(linear.bias.data) + + # Initialize LoRA parameters (matches the previous init_lora_weights defaults) + if not module.u.is_meta: + nn.init.normal_(module.u, mean=0.0, std=0.02) + nn.init.zeros_(module.v) + nn.init.ones_(module.sigma) + return module def forward(self, input): From a1a8fab05cfe63fc14a6982aa04136a7cd7198d9 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 17 Jul 2026 01:56:11 -0500 Subject: [PATCH 08/14] fix: propagate lora_rank to DecomposedLinear.Config --- alto/modifiers/lpt/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 849819a9..79879076 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -221,6 +221,7 @@ def is_match( in_features=config.in_features, out_features=config.out_features, bias=config.bias, + lora_rank=config.lora_rank, param_init=config.param_init | DecomposedLinear._EXTRA_INIT, ) setattr(parent, attr, new_config) From 536c13db487847829a21ef41010e0fb1d5a2cb49 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 19 Jul 2026 21:31:48 -0500 Subject: [PATCH 09/14] fix: checkpoint manager failed to access untyped storage if tensor wrapper applied --- alto/kernels/dispatch/tensor.py | 3 +++ alto/modifiers/lpt/base.py | 1 - 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/alto/kernels/dispatch/tensor.py b/alto/kernels/dispatch/tensor.py index 40868afd..430ba838 100644 --- a/alto/kernels/dispatch/tensor.py +++ b/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, diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 79879076..863655be 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -214,7 +214,6 @@ def is_match( if self.lora_rank > 0: resolved_targets = list(self.resolved_config.values()) for _fqn, config, parent, attr in model_config.traverse(Linear.Config): - #raise ValueError(f"[{_fqn}] cfg={config},\nparent={parent},\nattr={attr}") for target in resolved_targets: if is_match(_fqn, "Linear", target, self.ignore): new_config = DecomposedLinear.Config( From db5fcf3ae52e8c2b20d06c11a341c2e6908421e1 Mon Sep 17 00:00:00 2001 From: Hann Wang Date: Mon, 20 Jul 2026 14:42:41 +0800 Subject: [PATCH 10/14] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- alto/models/patcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alto/models/patcher.py b/alto/models/patcher.py index 9497e543..3059b0f4 100644 --- a/alto/models/patcher.py +++ b/alto/models/patcher.py @@ -129,4 +129,4 @@ def apply_rotary_emb_complex( ).contiguous() return original_apply_rotary_emb_complex(xq, xk, rope_cache) - ComplexRoPE.apply_rotary_emb = apply_rotary_emb_complex + ComplexRoPE.apply_rotary_emb = classmethod(apply_rotary_emb_complex) From dfc194a91f20878f5122c8ec9b19f6ec35acf290 Mon Sep 17 00:00:00 2001 From: Hann Wang Date: Mon, 20 Jul 2026 14:44:13 +0800 Subject: [PATCH 11/14] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- alto/modifiers/lpt/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 863655be..43e2f941 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -220,7 +220,7 @@ def is_match( in_features=config.in_features, out_features=config.out_features, bias=config.bias, - lora_rank=config.lora_rank, + lora_rank=self.lora_rank, param_init=config.param_init | DecomposedLinear._EXTRA_INIT, ) setattr(parent, attr, new_config) From 24ed297a546e6d14f3b904d73f90f39923facaf2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 24 Jul 2026 03:50:27 +0000 Subject: [PATCH 12/14] fx: missing trainer for de-osc --- 3rdparty/torchtitan | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan index cfa631dd..8bebda9e 160000 --- a/3rdparty/torchtitan +++ b/3rdparty/torchtitan @@ -1 +1 @@ -Subproject commit cfa631dd08919107867a65694da32b3faf4fd46d +Subproject commit 8bebda9ec7a98e097103956c1daf3a65128cf85e From 5e20087659c07ab9236da291fdb1c4f1fd5b1aae Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 27 Jul 2026 02:14:18 +0000 Subject: [PATCH 13/14] [WIP] feat: madam optimizer --- 3rdparty/torchtitan | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/torchtitan b/3rdparty/torchtitan index 8bebda9e..62d58153 160000 --- a/3rdparty/torchtitan +++ b/3rdparty/torchtitan @@ -1 +1 @@ -Subproject commit 8bebda9ec7a98e097103956c1daf3a65128cf85e +Subproject commit 62d58153183939fae36266ff3936955ccaa25375 From ec147ef696a6a8e03c4e4c448ecc1da2e106a472 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 28 Jul 2026 07:40:29 +0000 Subject: [PATCH 14/14] test: decompsoed linear with svd init --- tests/unittest/nn/test_decomposed_linear.py | 95 ++++++++++++++++++++- tests/unittest/nn/utils.py | 1 + 2 files changed, 94 insertions(+), 2 deletions(-) create mode 120000 tests/unittest/nn/utils.py diff --git a/tests/unittest/nn/test_decomposed_linear.py b/tests/unittest/nn/test_decomposed_linear.py index 191540fa..b0baa3e5 100644 --- a/tests/unittest/nn/test_decomposed_linear.py +++ b/tests/unittest/nn/test_decomposed_linear.py @@ -3,12 +3,16 @@ # SPDX-License-Identifier: MIT import pytest +from tabulate import tabulate import torch import torch.nn as nn import torch.nn.functional as F from alto.nn import DecomposedLinear +from alto.kernels.fp4.mxfp4.mxfp_quantization import convert_to_mxfp4, convert_from_mxfp4 from alto.kernels.dispatch import TrainingOpConfig, swap_params +from utils import prepare_data, calc_cossim, calc_snr + @pytest.mark.parametrize("in_features", [128, 256, 512]) @pytest.mark.parametrize("out_features", [128, 256, 512]) @@ -28,7 +32,8 @@ def test_decomposed_linear(in_features, out_features, bias, lora_rank): decomposed_linear.sigma.data.normal_(mean=0, std=STD) y_decomposed = decomposed_linear(x) - combined_weight = decomposed_linear.weight + decomposed_linear.u.T @ torch.diag(decomposed_linear.sigma) @ decomposed_linear.v.T + combined_weight = decomposed_linear.weight + decomposed_linear.u.T @ torch.diag( + decomposed_linear.sigma) @ decomposed_linear.v.T y_ref = F.linear(x, combined_weight, bias=decomposed_linear.bias) assert torch.allclose(y_ref, y_decomposed, atol=1e-4, rtol=1e-4) @@ -41,7 +46,9 @@ def test_decomposed_linear(in_features, out_features, bias, lora_rank): @pytest.mark.parametrize("precision", ["mxfp4"]) def test_decomposed_linear_quantization(in_features, out_features, bias, lora_rank, precision): STD = 0.1 - decomposed_linear = DecomposedLinear(DecomposedLinear.Config(in_features=in_features, out_features=out_features, bias=bias, lora_rank=lora_rank)).to("cuda") + decomposed_linear = DecomposedLinear( + DecomposedLinear.Config(in_features=in_features, out_features=out_features, bias=bias, + lora_rank=lora_rank)).to("cuda") decomposed_linear.weight.data.normal_(mean=0, std=STD) if bias: decomposed_linear.bias.data.normal_(mean=0, std=STD) @@ -71,3 +78,87 @@ def test_decomposed_linear_quantization(in_features, out_features, bias, lora_ra cossim = torch.nn.functional.cosine_similarity(y_ref.flatten(), y.flatten(), dim=0) print(f"snr={snr}, cossim={cossim}, max_diff={max_diff}, mean_diff={mean_diff}") assert snr > 10, f"SNR too low: {snr}" + + +@pytest.mark.parametrize("shape", [(4, 1024, 1024, 2048)]) +@pytest.mark.parametrize("data_type", [torch.float32]) +def test_decomposed_linear_with_svd(shape, data_type): + B, M, N, K = shape + inputs = prepare_data((B, M, K), data_type).requires_grad_(True) + weights = prepare_data((N, K), data_type).requires_grad_(True) + + # Create a target for gradient computation + target = prepare_data((B, M, N), data_type) + + # PyTorch reference implementation + linear = torch.nn.Linear(K, N, bias=False) + linear.weight = torch.nn.Parameter(weights) + outputs_ref = linear(inputs) + + # Compute loss and gradients with PyTorch + loss_ref = torch.nn.functional.mse_loss(outputs_ref, target) + + loss_ref.backward() + grad_inputs_ref = inputs.grad.clone() + grad_weights_ref = linear.weight.grad.clone() + + # Reset gradients + inputs.grad.zero_() + linear.weight.grad.zero_() + + decomposed_linear = DecomposedLinear.from_linear(linear, lora_rank=32) + + with torch.no_grad(): + w_fp4, w_scales = convert_to_mxfp4(weights, axis=-1, is_2d_block=True) + w_qdq = convert_from_mxfp4( + w_fp4, + w_scales, + output_dtype=data_type, + axis=-1, + is_2d_block=True, + ) + decomposed_linear.weight.data.copy_(w_qdq) + + delta = weights - w_qdq + u, sigma, v = torch.svd_lowrank(delta, q=32) + + decomposed_linear.u.data.copy_(u.T) + decomposed_linear.sigma.data.copy_(sigma) + decomposed_linear.v.data.copy_(v) + + # decomposed_linear.u.data.normal_(mean=0, std=STD) + # decomposed_linear.v.data.normal_(mean=0, std=STD) + # decomposed_linear.sigma.data.normal_(mean=0, std=STD) + + quant_op_config = TrainingOpConfig( + precision="mxfp4", + use_2dblock_x=False, + use_2dblock_w=True, + use_hadamard=True, + use_sr_grad=False, + use_dge=False, + ) + swap_params(decomposed_linear, config=quant_op_config, target_parameter_name="weight") + swap_params(decomposed_linear, config=quant_op_config, target_parameter_name="u") + swap_params(decomposed_linear, config=quant_op_config, target_parameter_name="v") + outputs = decomposed_linear(inputs) + + loss = torch.nn.functional.mse_loss(outputs, target) + loss.backward() + + output_snr = calc_snr(outputs, outputs_ref) + output_sim = calc_cossim(outputs, outputs_ref) + dx_snr = calc_snr(inputs.grad, grad_inputs_ref) + dx_sim = calc_cossim(inputs.grad, grad_inputs_ref) + dw_snr = calc_snr(linear.weight.grad, grad_weights_ref) + dw_sim = calc_cossim(linear.weight.grad, grad_weights_ref) + + print() + print( + tabulate([ + ["O", output_snr, output_sim], + ["dX", dx_snr, dx_sim], + ["dW", dw_snr, dw_sim], + ], + headers=["Tensor", "SNR", "Cosine Sim"], + tablefmt="github")) diff --git a/tests/unittest/nn/utils.py b/tests/unittest/nn/utils.py new file mode 120000 index 00000000..67c780fc --- /dev/null +++ b/tests/unittest/nn/utils.py @@ -0,0 +1 @@ +../mxfp4/utils.py \ No newline at end of file