From 18a712b90496405ec81d30870d09c3bb59d4b512 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 28 May 2026 17:31:03 +0000 Subject: [PATCH 1/9] Revert "[`Revert`] FSDP+Dtensor refactor related changes (#46246)" This reverts commit 295cee3e1d0085fad27b3ff1da8aa03af8d0c565. --- .circleci/create_circleci_config.py | 12 +- src/transformers/__init__.py | 1 + src/transformers/configuration_utils.py | 13 + src/transformers/core_model_loading.py | 125 +- src/transformers/distributed/__init__.py | 14 + .../distributed/configuration_utils.py | 138 +- src/transformers/distributed/fsdp.py | 398 +++++ .../distributed/sharding_utils.py | 495 ++++++ .../distributed/tensor_parallel.py | 526 ++++++ src/transformers/distributed/utils.py | 299 ++++ src/transformers/generation/utils.py | 6 +- src/transformers/integrations/__init__.py | 13 - src/transformers/integrations/accelerate.py | 2 +- src/transformers/integrations/fsdp.py | 92 - src/transformers/integrations/moe.py | 2 +- src/transformers/integrations/mxfp4.py | 42 +- .../integrations/tensor_parallel.py | 1560 ----------------- src/transformers/modeling_utils.py | 371 ++-- .../models/afmoe/modeling_afmoe.py | 4 +- .../models/afmoe/modular_afmoe.py | 4 +- .../models/apertus/configuration_apertus.py | 29 +- .../models/apertus/modeling_apertus.py | 4 +- .../models/apertus/modular_apertus.py | 29 +- .../models/arcee/configuration_arcee.py | 24 +- .../models/arcee/modeling_arcee.py | 4 +- .../models/arcee/modular_arcee.py | 18 +- .../models/aria/configuration_aria.py | 25 +- src/transformers/models/aria/modeling_aria.py | 4 +- src/transformers/models/aria/modular_aria.py | 4 +- .../audioflamingo3/modeling_audioflamingo3.py | 1 + .../audioflamingo3/modular_audioflamingo3.py | 1 + .../models/bamba/modeling_bamba.py | 4 +- .../models/bitnet/modeling_bitnet.py | 2 + .../models/bitnet/modular_bitnet.py | 1 + .../models/cohere/configuration_cohere.py | 27 +- .../models/cohere/modeling_cohere.py | 4 +- .../models/cohere2/configuration_cohere2.py | 25 +- .../models/cohere2/modeling_cohere2.py | 4 +- .../models/cohere2/modular_cohere2.py | 25 +- .../cohere2_moe/configuration_cohere2_moe.py | 49 +- .../cohere2_moe/modeling_cohere2_moe.py | 4 +- src/transformers/models/csm/modeling_csm.py | 2 + src/transformers/models/csm/modular_csm.py | 1 + .../models/cwm/configuration_cwm.py | 25 +- src/transformers/models/cwm/modeling_cwm.py | 4 +- .../data2vec/modeling_data2vec_audio.py | 2 +- src/transformers/models/dbrx/modeling_dbrx.py | 4 +- src/transformers/models/dbrx/modular_dbrx.py | 4 +- .../deepseek_v2/configuration_deepseek_v2.py | 41 +- .../deepseek_v2/modeling_deepseek_v2.py | 4 +- .../models/deepseek_v2/modular_deepseek_v2.py | 35 +- .../deepseek_v3/configuration_deepseek_v3.py | 29 +- .../deepseek_v3/modeling_deepseek_v3.py | 4 +- .../deepseek_v4/configuration_deepseek_v4.py | 8 +- .../deepseek_v4/modeling_deepseek_v4.py | 4 +- src/transformers/models/dia/generation_dia.py | 2 +- .../models/diffllama/modeling_diffllama.py | 4 +- .../models/doge/configuration_doge.py | 18 +- src/transformers/models/doge/modeling_doge.py | 4 +- src/transformers/models/doge/modular_doge.py | 18 +- .../models/dots1/configuration_dots1.py | 19 +- .../models/dots1/modeling_dots1.py | 4 +- .../models/dots1/modular_dots1.py | 19 +- src/transformers/models/emu3/modeling_emu3.py | 4 +- .../models/ernie4_5/configuration_ernie4_5.py | 25 +- .../models/ernie4_5/modeling_ernie4_5.py | 4 +- .../configuration_ernie4_5_moe.py | 16 +- .../ernie4_5_moe/modeling_ernie4_5_moe.py | 4 +- .../configuration_ernie4_5_vl_moe.py | 16 +- .../modular_ernie4_5_vl_moe.py | 10 +- .../models/eurobert/configuration_eurobert.py | 25 +- .../models/eurobert/modeling_eurobert.py | 4 +- .../models/eurobert/modular_eurobert.py | 4 +- .../models/exaone4/configuration_exaone4.py | 29 +- .../models/exaone4/modeling_exaone4.py | 4 +- .../models/exaone4/modular_exaone4.py | 29 +- .../exaone_moe/configuration_exaone_moe.py | 14 +- .../models/exaone_moe/modeling_exaone_moe.py | 4 +- .../models/exaone_moe/modular_exaone_moe.py | 8 + .../models/falcon_h1/modeling_falcon_h1.py | 4 +- .../flex_olmo/configuration_flex_olmo.py | 17 +- .../models/flex_olmo/modeling_flex_olmo.py | 4 +- .../models/flex_olmo/modular_flex_olmo.py | 17 +- .../models/gemma/configuration_gemma.py | 25 +- .../models/gemma/modeling_gemma.py | 4 +- .../models/gemma/modular_gemma.py | 25 +- .../models/gemma2/configuration_gemma2.py | 25 +- .../models/gemma2/modeling_gemma2.py | 4 +- .../models/gemma2/modular_gemma2.py | 25 +- .../models/gemma3/configuration_gemma3.py | 30 +- .../models/gemma3/modeling_gemma3.py | 4 +- .../models/gemma3/modular_gemma3.py | 30 +- .../models/gemma3n/configuration_gemma3n.py | 21 +- .../models/gemma3n/modeling_gemma3n.py | 4 +- .../models/gemma3n/modular_gemma3n.py | 21 +- .../models/gemma4/configuration_gemma4.py | 33 +- .../models/gemma4/modeling_gemma4.py | 4 +- .../modeling_gemma4_assistant.py | 1 + .../models/glm/configuration_glm.py | 26 +- src/transformers/models/glm/modeling_glm.py | 4 +- .../models/glm4/configuration_glm4.py | 26 +- src/transformers/models/glm4/modeling_glm4.py | 4 +- .../models/glm4_moe/configuration_glm4_moe.py | 17 +- .../models/glm4_moe/modeling_glm4_moe.py | 4 +- .../models/glm4_moe/modular_glm4_moe.py | 17 +- .../configuration_glm4_moe_lite.py | 16 +- .../glm4_moe_lite/modeling_glm4_moe_lite.py | 4 +- .../glm4_moe_lite/modular_glm4_moe_lite.py | 16 +- .../models/glm4v/configuration_glm4v.py | 13 +- .../models/glm4v/modular_glm4v.py | 13 +- .../glm4v_moe/configuration_glm4v_moe.py | 12 +- .../models/glm4v_moe/modular_glm4v_moe.py | 11 +- .../glm_image/configuration_glm_image.py | 13 +- .../models/glm_image/modular_glm_image.py | 6 + .../glm_moe_dsa/configuration_glm_moe_dsa.py | 17 +- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 4 +- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 17 +- .../models/glm_ocr/configuration_glm_ocr.py | 13 +- .../models/glm_ocr/modular_glm_ocr.py | 6 + .../models/glmasr/modeling_glmasr.py | 1 + .../models/gpt_neox/configuration_gpt_neox.py | 19 +- .../models/gpt_neox/modeling_gpt_neox.py | 2 +- .../models/gpt_neox/modular_gpt_neox.py | 2 +- .../models/gpt_oss/configuration_gpt_oss.py | 7 +- .../models/gpt_oss/modeling_gpt_oss.py | 4 +- .../models/granite/configuration_granite.py | 25 +- .../models/granite/modeling_granite.py | 4 +- .../configuration_granite4_vision.py | 25 +- .../models/granitemoe/modeling_granitemoe.py | 4 +- .../modeling_granitemoehybrid.py | 4 +- .../modeling_granitemoeshared.py | 4 +- .../models/helium/configuration_helium.py | 25 +- .../models/helium/modeling_helium.py | 4 +- .../configuration_higgs_audio_v2.py | 25 +- .../models/hrm_text/configuration_hrm_text.py | 21 + .../models/hrm_text/modeling_hrm_text.py | 4 +- .../models/hubert/modeling_hubert.py | 2 +- .../modeling_hunyuan_v1_dense.py | 4 +- .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 4 +- .../models/hy_v3/configuration_hy_v3.py | 20 +- .../models/hy_v3/modeling_hy_v3.py | 4 +- .../models/hy_v3/modular_hy_v3.py | 20 +- .../hyperclovax/configuration_hyperclovax.py | 25 +- .../hyperclovax/modeling_hyperclovax.py | 4 +- .../models/jais2/configuration_jais2.py | 24 +- .../models/jais2/modeling_jais2.py | 4 +- .../models/jais2/modular_jais2.py | 18 +- .../models/jamba/modeling_jamba.py | 4 +- .../modeling_kyutai_speech_to_text.py | 4 +- .../models/laguna/configuration_laguna.py | 20 +- .../models/laguna/modeling_laguna.py | 4 +- .../models/laguna/modular_laguna.py | 14 +- src/transformers/models/lfm2/modeling_lfm2.py | 4 +- .../models/lfm2_moe/modeling_lfm2_moe.py | 4 +- .../models/llama/configuration_llama.py | 25 +- .../models/llama/modeling_llama.py | 4 +- .../models/llama4/configuration_llama4.py | 20 +- .../models/llama4/modeling_llama4.py | 4 +- .../configuration_longcat_flash.py | 15 +- .../longcat_flash/modeling_longcat_flash.py | 4 +- .../models/minimax/configuration_minimax.py | 13 +- .../models/minimax/modeling_minimax.py | 4 +- .../models/minimax/modular_minimax.py | 13 +- .../minimax_m2/configuration_minimax_m2.py | 32 +- .../models/minimax_m2/modeling_minimax_m2.py | 4 +- .../models/minimax_m2/modular_minimax_m2.py | 32 +- .../ministral/configuration_ministral.py | 25 +- .../models/ministral/modeling_ministral.py | 4 +- .../ministral3/configuration_ministral3.py | 25 +- .../models/ministral3/modeling_ministral3.py | 4 +- .../models/mistral/configuration_mistral.py | 25 +- .../models/mistral/modeling_mistral.py | 4 +- .../models/mistral4/configuration_mistral4.py | 14 +- .../models/mistral4/modeling_mistral4.py | 4 +- .../models/mixtral/configuration_mixtral.py | 29 +- .../models/mixtral/modeling_mixtral.py | 4 +- .../models/mixtral/modular_mixtral.py | 1 + .../musicflamingo/modeling_musicflamingo.py | 1 + .../models/nanochat/configuration_nanochat.py | 5 +- .../models/nanochat/modeling_nanochat.py | 4 +- .../models/nanochat/modular_nanochat.py | 4 +- .../models/nllb_moe/modeling_nllb_moe.py | 2 +- .../models/olmo/configuration_olmo.py | 25 +- src/transformers/models/olmo/modeling_olmo.py | 4 +- .../models/olmo2/configuration_olmo2.py | 33 +- .../models/olmo2/modeling_olmo2.py | 4 +- .../models/olmo2/modular_olmo2.py | 27 +- .../models/olmo3/configuration_olmo3.py | 33 +- .../models/olmo3/modeling_olmo3.py | 4 +- .../models/olmo3/modular_olmo3.py | 27 +- .../olmo_hybrid/configuration_olmo_hybrid.py | 17 +- .../olmo_hybrid/modeling_olmo_hybrid.py | 4 +- .../models/olmo_hybrid/modular_olmo_hybrid.py | 11 +- .../models/olmoe/configuration_olmoe.py | 31 +- .../models/olmoe/modeling_olmoe.py | 4 +- .../configuration_openai_privacy_filter.py | 9 +- .../modular_openai_privacy_filter.py | 7 + .../configuration_paddleocr_vl.py | 25 +- .../models/phi/configuration_phi.py | 28 +- src/transformers/models/phi/modeling_phi.py | 4 +- .../models/phi3/configuration_phi3.py | 26 +- src/transformers/models/phi3/modeling_phi3.py | 4 +- .../configuration_phi4_multimodal.py | 26 +- .../modeling_phi4_multimodal.py | 4 +- .../models/phimoe/configuration_phimoe.py | 5 + .../models/phimoe/modeling_phimoe.py | 4 +- src/transformers/models/pi0/modeling_pi0.py | 2 +- src/transformers/models/pi0/modular_pi0.py | 2 +- .../models/qwen2/configuration_qwen2.py | 25 +- .../models/qwen2/modeling_qwen2.py | 4 +- .../configuration_qwen2_5_omni.py | 11 +- .../qwen2_5_omni/modular_qwen2_5_omni.py | 11 +- .../qwen2_5_vl/configuration_qwen2_5_vl.py | 10 +- .../qwen2_moe/configuration_qwen2_moe.py | 10 +- .../models/qwen2_moe/modeling_qwen2_moe.py | 4 +- .../models/qwen2_moe/modular_qwen2_moe.py | 4 +- .../models/qwen2_vl/configuration_qwen2_vl.py | 10 +- .../models/qwen3/configuration_qwen3.py | 43 +- .../models/qwen3/modeling_qwen3.py | 4 +- .../models/qwen3/modular_qwen3.py | 2 + .../models/qwen3_5/configuration_qwen3_5.py | 12 +- .../models/qwen3_5/modeling_qwen3_5.py | 4 +- .../models/qwen3_5/modular_qwen3_5.py | 6 +- .../qwen3_5_moe/configuration_qwen3_5_moe.py | 16 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 8 +- .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 14 +- .../qwen3_moe/configuration_qwen3_moe.py | 36 +- .../models/qwen3_moe/modeling_qwen3_moe.py | 4 +- .../qwen3_next/configuration_qwen3_next.py | 18 +- .../models/qwen3_next/modeling_qwen3_next.py | 4 +- .../configuration_qwen3_omni_moe.py | 92 +- .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 8 +- .../qwen3_omni_moe/modular_qwen3_omni_moe.py | 17 +- .../configuration_qwen3_vl_moe.py | 30 +- .../qwen3_vl_moe/modular_qwen3_vl_moe.py | 6 +- .../seamless_m4t/modeling_seamless_m4t.py | 2 +- .../modeling_seamless_m4t_v2.py | 2 +- .../models/seed_oss/configuration_seed_oss.py | 25 +- .../models/seed_oss/modeling_seed_oss.py | 4 +- src/transformers/models/sew/modeling_sew.py | 2 +- src/transformers/models/sew/modular_sew.py | 2 +- .../models/smollm3/configuration_smollm3.py | 25 +- .../models/smollm3/modeling_smollm3.py | 4 +- .../models/smollm3/modular_smollm3.py | 25 +- .../solar_open/configuration_solar_open.py | 26 +- .../models/solar_open/modeling_solar_open.py | 4 +- .../models/solar_open/modular_solar_open.py | 25 +- .../models/speecht5/modeling_speecht5.py | 2 +- .../starcoder2/configuration_starcoder2.py | 24 +- .../models/starcoder2/modeling_starcoder2.py | 4 +- .../models/t5gemma/configuration_t5gemma.py | 25 +- .../models/t5gemma/modeling_t5gemma.py | 2 +- .../models/t5gemma/modular_t5gemma.py | 8 +- .../models/t5gemma2/configuration_t5gemma2.py | 60 +- .../models/t5gemma2/modeling_t5gemma2.py | 2 +- .../models/t5gemma2/modular_t5gemma2.py | 9 +- .../models/unispeech/modeling_unispeech.py | 2 +- .../unispeech_sat/modeling_unispeech_sat.py | 2 +- .../vaultgemma/configuration_vaultgemma.py | 25 +- .../models/vaultgemma/modeling_vaultgemma.py | 4 +- .../models/vaultgemma/modular_vaultgemma.py | 6 + src/transformers/models/vits/modeling_vits.py | 2 +- .../configuration_voxtral_realtime.py | 10 +- .../models/wav2vec2/modeling_wav2vec2.py | 2 +- .../wav2vec2_bert/modeling_wav2vec2_bert.py | 2 +- .../wav2vec2_bert/modular_wav2vec2_bert.py | 2 +- .../modeling_wav2vec2_conformer.py | 2 +- .../modular_wav2vec2_conformer.py | 2 +- .../models/wavlm/modeling_wavlm.py | 2 +- .../models/wavlm/modular_wavlm.py | 2 +- .../models/youtu/configuration_youtu.py | 24 +- .../models/youtu/modeling_youtu.py | 4 +- .../models/youtu/modular_youtu.py | 3 +- src/transformers/testing_utils.py | 118 +- src/transformers/trainer.py | 9 +- src/transformers/trainer_seq2seq.py | 2 +- tests/causal_lm_tester.py | 8 +- tests/tensor_parallel/test_tensor_parallel.py | 309 ---- tests/test_distributed_config.py | 93 + tests/test_fsdp_mixin.py | 784 +++++++++ tests/test_modeling_common.py | 2 +- tests/test_tensor_parallel_mixin.py | 136 +- tests/utils/test_core_model_loading.py | 115 +- .../utils/test_distributed_sharding_utils.py | 462 +++++ tests/utils/test_modeling_utils.py | 38 + utils/tests_fetcher.py | 1 + 286 files changed, 6332 insertions(+), 3011 deletions(-) create mode 100644 src/transformers/distributed/fsdp.py create mode 100644 src/transformers/distributed/sharding_utils.py create mode 100644 src/transformers/distributed/tensor_parallel.py create mode 100644 src/transformers/distributed/utils.py delete mode 100644 src/transformers/integrations/fsdp.py delete mode 100644 src/transformers/integrations/tensor_parallel.py create mode 100644 tests/test_distributed_config.py create mode 100644 tests/test_fsdp_mixin.py create mode 100644 tests/utils/test_distributed_sharding_utils.py diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py index 26f254d141b7..54c69567c159 100644 --- a/.circleci/create_circleci_config.py +++ b/.circleci/create_circleci_config.py @@ -398,6 +398,15 @@ def job_name(self): parallelism=6, ) +fsdp_ci_job = CircleCIJob( + "fsdp_ci", + additional_env={"RUN_FSDP_TESTS": True}, + docker_image=[{"image": "huggingface/transformers-torch-light"}], + install_steps=["uv pip install .", "uv pip install torchao"], + marker="is_fsdp_test", + parallelism=6, +) + # We also include a `dummy.py` file in the files to be doc-tested to prevent edge case failure. Otherwise, the pytest # hangs forever during test collection while showing `collecting 0 items / 21 errors`. (To see this, we have to remove # the bash output redirection.) @@ -429,7 +438,8 @@ def job_name(self): DOC_TESTS = [doc_test_job] TRAINING_CI_TESTS = [training_ci_job] TENSOR_PARALLEL_CI_TESTS = [tensor_parallel_ci_job] -ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS # fmt: skip +FSDP_CI_TESTS = [fsdp_ci_job] +ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS + FSDP_CI_TESTS # fmt: skip def create_circleci_config(folder=None): diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 9b05c92dabf1..1d629f361d82 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -104,6 +104,7 @@ "debug_utils": [], "dependency_versions_check": [], "dependency_versions_table": [], + "distributed": [], "dynamic_module_utils": [], "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"], "feature_extraction_utils": ["BatchFeature", "FeatureExtractionMixin"], diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 04cc14de45c1..f5b8d04da4ee 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -147,6 +147,13 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): naming of attributes. - **base_model_tp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a tensor parallel plan applied to the sub-module when `model.tensor_parallel` is called. + - **base_model_sp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a sequence + parallel plan, used in place of `base_model_tp_plan` when `distributed_config.enable_sequence_parallel` is set. + Same key/value shape as the TP plan; values are style names registered in `ALL_PARALLEL_STYLES` + (e.g. `"vocab_reduce_scatter"`, `"rowwise_reduce_scatter"`, `"activation"`, `"module_allgather"`). + - **base_model_fsdp_plan** (`dict[Any, str]`) -- A dict that maps sub-modules of a base model to an FSDP2 + sharding strategy (e.g. `"free_full_weight"` / `"keep_full_weight"`). Keys can be wildcard module paths + (e.g. `"layers.*"`) or tuples of paths (grouped into a single `fully_shard` call). - **base_model_pp_plan** (`dict[str, tuple[list[str]]]`) -- A dict that maps child-modules of a base model to a pipeline parallel plan that enables users to place the child-module on the appropriate device. @@ -218,6 +225,8 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): keys_to_ignore_at_inference: ClassVar[list[str]] = [] attribute_map: ClassVar[dict[str, str]] = {} base_model_tp_plan: ClassVar[dict[str, Any] | None] = None + base_model_sp_plan: ClassVar[dict[str, Any] | None] = None + base_model_fsdp_plan: ClassVar[dict[Any, str] | None] = None base_model_pp_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None base_model_ep_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None _auto_class: ClassVar[str | None] = None @@ -1020,6 +1029,9 @@ def to_dict(self) -> dict[str, Any]: # Pop "kwargs" since they are unpacked and set in the post init output.pop("kwargs", None) + if "distributed_config" in output and hasattr(output["distributed_config"], "to_dict"): + output["distributed_config"] = output["distributed_config"].to_dict() + def to_list(value): if isinstance(value, tuple): value = [to_list(item) for item in value] @@ -1165,6 +1177,7 @@ def _remove_keys_not_serialized(self, d: dict[str, Any]) -> None: "_experts_implementation_internal", "ignore_keys_at_rope_validation", "base_model_tp_plan", + "base_model_sp_plan", "base_model_pp_plan", ]: d.pop(key_to_remove, None) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 53540e53ca61..271e1898d4db 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -30,21 +30,21 @@ import torch +from .distributed.sharding_utils import DtensorShardOperation, _dtensor_from_local_like from .integrations.accelerate import get_device, offload_weight -from .integrations.tensor_parallel import ALL_PARALLEL_STYLES from .utils import is_env_variable_true from .utils.loading_report import LoadStateDictInfo from .utils.logging import get_logger, tqdm _torch_distributed_available = torch.distributed.is_available() +if _torch_distributed_available: + from torch.distributed.tensor import DTensor if TYPE_CHECKING: - from .integrations.tensor_parallel import TensorParallelLayer from .modeling_utils import LoadStateDictConfig, PreTrainedModel from .quantizers import HfQuantizer - logger = get_logger(__name__) @@ -392,7 +392,7 @@ def __init__(self): def _apply(self, tensor: torch.Tensor) -> torch.Tensor: dim1, dim2 = tensor.shape - n_heads = self.config.getattr("num_attention_heads", 1) + n_heads = getattr(self.config, "num_attention_heads", 1) tensor = tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) tensor = tensor.transpose(1, 2).reshape(dim1, dim2) @@ -408,11 +408,10 @@ def convert( **kwargs, ) -> dict[str, list[torch.Tensor]]: self.config = config - output: dict[str, list[torch.Tensor]] = {} + output = {} for key, tensors in input_dict.items(): - if len(tensors) != 1: - raise ValueError("PermuteForRope expects a single tensor per key.") - output[key] = [self._apply(tensors[0])] + tensor = tensors[0] if isinstance(tensors, list) else tensors + output[key] = self._apply(tensor) return output @@ -614,7 +613,7 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list self._original_target_patterns = self.target_patterns.copy() # Init fields that will be used during conversion - self.distributed_operation: TensorParallelLayer | None = None + self.distributed_operation: Any = None self.quantization_operation: ConversionOps | None = None self.collected_tensors: dict[str, list[Future]] = defaultdict(list) self.layer_targets: dict[str, set[str]] = defaultdict(set) @@ -772,7 +771,9 @@ def reverse_transform(self) -> WeightTransform: kwargs["operations"] = [op.reverse_op for op in self.operations[::-1]] reverse_transform = self.__class__( - source_patterns=self._original_target_patterns, target_patterns=self._original_source_patterns, **kwargs + source_patterns=self._original_target_patterns, + target_patterns=self._original_source_patterns, + **kwargs, ) reverse_transform.scope_prefix = self.scope_prefix reverse_transform.base_model_prefix = self.base_model_prefix @@ -799,6 +800,8 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: # Sync loading elif callable(tensors[0]): tensors = [func() for func in tensors] + # Some may be None for some distributed setups + tensors = [tensor for tensor in tensors if tensor is not None] # Add them to the new dictionary collected_tensors[key] = tensors @@ -1001,29 +1004,20 @@ def spawn_materialize( tensor: torch.Tensor, device=None, dtype=None, + sharding_op: DtensorShardOperation | None = None, + tensor_idx: int | None = None, ) -> Future | Callable: - """Materialize a tensor from file asynchronously if `thread_pool` is provided, or return a Callable that will - load the tensor synchronously when called.""" - - def _job(): - return _materialize_copy(tensor, device, dtype) - - if thread_pool is not None: - return thread_pool.submit(_job) - else: - # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu - # memory during Conversion - return _job - + """Materialize (and optionally shard) a tensor, asynchronously if a thread pool is provided. -def spawn_tp_materialize( - thread_pool: ThreadPoolExecutor | None, tensor: torch.Tensor, sharding_method, tensor_idx, device=None, dtype=None -) -> Future | Callable: - """Materialize and shard a tensor (according to the TP-plan) from file asynchronously if `thread_pool` is provided, or - return a Callable that will load the tensor synchronously when called.""" + When ``sharding_op`` is given the tensor is sharded according to the DTensor + placement strategy; otherwise it is simply copied to *device*/*dtype*. + Without a thread pool a deferred callable is returned instead of a Future. + """ def _job(): - return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype) + if sharding_op is not None: + return sharding_op.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype) + return _materialize_copy(tensor, device, dtype) if thread_pool is not None: return thread_pool.submit(_job) @@ -1096,12 +1090,12 @@ def _format_op_name(curr_op: list[ConversionOps] | ConversionOps | None) -> str raise SkipParameters() +@torch.no_grad() def set_param_for_module( model: PreTrainedModel, target_name: str, param_value: torch.Tensor, loading_info: LoadStateDictInfo, - distributed_operation: TensorParallelLayer | None, hf_quantizer: HfQuantizer, ): module_path, _, param_name = target_name.rpartition(".") @@ -1116,27 +1110,25 @@ def set_param_for_module( if ref is None: loading_info.unexpected_keys.add(target_name) else: - if not isinstance(param_value, torch.nn.Parameter): + if not isinstance(param_value, torch.nn.Parameter) and not isinstance(ref, DTensor): if param_name not in module_obj._buffers: param_value = torch.nn.Parameter(param_value, requires_grad=param_value.is_floating_point()) # Remove from missing keys (it's either mismatched, or all good) loading_info.missing_keys.discard(target_name) - # Determine expected shape: for TP, use sharded shape; otherwise, use full shape - if distributed_operation is not None: - expected_shape = torch.Size(distributed_operation.get_expected_sharded_shape(ref.shape)) - else: - expected_shape = ref.shape + expected_shape = ref._local_tensor.shape if isinstance(ref, DTensor) else ref.shape if ref is not None and param_value.shape != expected_shape and hf_quantizer is None: loading_info.mismatched_keys.add((target_name, param_value.shape, expected_shape)) else: + if isinstance(ref, DTensor): + local_param = param_value.detach() if isinstance(param_value, torch.nn.Parameter) else param_value + dtensor_param = _dtensor_from_local_like(local_param, ref) + param_value = torch.nn.Parameter(dtensor_param, requires_grad=ref.requires_grad) # super important otherwise _init_weight will re-init the param param_value._is_hf_initialized = True setattr(module_obj, param_name, param_value) - if distributed_operation is not None: - distributed_operation.update_module_attributes(module_obj) def offload_and_maybe_resave_param( @@ -1321,7 +1313,6 @@ def convert_and_load_state_dict_in_model( device_map = load_config.device_map or {"": "cpu"} hf_quantizer = load_config.hf_quantizer dtype = load_config.dtype - device_mesh = load_config.device_mesh disk_offload_folder = load_config.disk_offload_folder offload_buffers = load_config.offload_buffers dtype_plan = load_config.dtype_plan or {} @@ -1356,10 +1347,6 @@ def convert_and_load_state_dict_in_model( converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)] param_name_to_load: dict[str, WeightRenaming | WeightConverter] = {} - # build '(?P.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'} - # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched. - if tp_plan != {}: - tp_plan_alt, tp_plan_by_group_name, _ = build_glob_alternation(list(tp_plan.keys())) if dtype_plan != {}: dtype_policy_alt, dtype_policy_by_group_name, _ = build_glob_alternation(list(dtype_plan.keys())) @@ -1423,32 +1410,23 @@ def convert_and_load_state_dict_in_model( elif empty_param is not None and empty_param.dtype != _dtype: _dtype = empty_param.dtype # usually correct when initializing - # 4. Handle TP sharding or device_map placement - future_or_tensor = None - if device_mesh and tp_plan: - if matched_tp_pattern := tp_plan_alt.search(renamed_key): - matched_tp_pattern = tp_plan_by_group_name[matched_tp_pattern.lastgroup] - if getattr(mapping, "distributed_operation", None) is None: - tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]].__class__ - mapping.distributed_operation = tp_layer( - device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone() - ) - shard_index = ( - len(mapping.collected_tensors.get(source_pattern, [])) - if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) - else None - ) - future_or_tensor = spawn_tp_materialize( - thread_pool, - tensor, - mapping.distributed_operation, - shard_index, - device_map[""], - _dtype, - ) - - if future_or_tensor is None: - param_device = get_device(device_map, renamed_key, valid_torch_device=True) + # 4. Materialize tensor — shard-on-read for DTensor params, plain copy otherwise + param_device = get_device(device_map, renamed_key, valid_torch_device=True) + if isinstance(empty_param, DTensor): + tensor_idx = ( + len(mapping.collected_tensors.get(source_pattern, [])) + if isinstance(mapping, WeightConverter) and isinstance(mapping.operations[0], MergeModulelist) + else None + ) + future_or_tensor = spawn_materialize( + thread_pool, + tensor, + param_device, + _dtype, + sharding_op=DtensorShardOperation(empty_param), + tensor_idx=tensor_idx, + ) + else: future_or_tensor = spawn_materialize(thread_pool, tensor, param_device, _dtype) mapping.add_tensor(renamed_key, original_key, source_pattern, future_or_tensor) @@ -1478,14 +1456,7 @@ def convert_and_load_state_dict_in_model( target_name, param, loading_info, disk_offload_folder, disk_offload_index, mapping ) else: - set_param_for_module( - model, - target_name, - param, - loading_info, - mapping.distributed_operation, - hf_quantizer, - ) + set_param_for_module(model, target_name, param, loading_info, hf_quantizer) # Cleanup all the tensors that were gathered before next iteration del realized_value diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index ba6db8358d2b..09f81832c8e0 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -19,6 +19,14 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], + "sharding_utils": [], + "tensor_parallel": [ + "ALL_PARALLEL_STYLES", + "apply_tensor_parallel", + "verify_tp_plan", + ], + "utils": [], } @@ -26,6 +34,12 @@ from .configuration_utils import ( DistributedConfig, ) + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module + from .tensor_parallel import ( + ALL_PARALLEL_STYLES, + apply_tensor_parallel, + verify_tp_plan, + ) else: import sys diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 7726d9f3290d..15c6a9e367f2 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -12,99 +12,75 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import json import os -from dataclasses import dataclass -from typing import Any +from dataclasses import asdict, dataclass + +from ..utils import is_torch_available + + +if is_torch_available(): + import torch @dataclass class DistributedConfig: """ - Base class for distributed configs + Configuration for native distributed training (FSDP2 + TP). + + Args: + tp_size (`int`, *optional*): + Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1. + tp_plan (`dict`, *optional*): + Tensor parallel sharding plan. Leave as `None` to use the model's default plan. + fsdp_size (`int`, *optional*): + Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. + fsdp_plan (`dict`, *optional*): + FSDP wrapping plan. Leave as `None` to wrap each transformer layer + root. """ + tp_size: int | None = None + tp_plan: dict[str, str] | None = None + enable_sequence_parallel: bool = False enable_expert_parallel: bool = False - # TODO: add tp_plan, pp_plan, device_mesh etc.. + fsdp_size: int | None = None + fsdp_plan: dict | None = None + + def __post_init__(self): + if self.tp_size is None and self.fsdp_size is None: + return + + if self.tp_size is None: + self.tp_size = 1 + if self.fsdp_size is None: + self.fsdp_size = 1 + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + if self.tp_size * self.fsdp_size != world_size: + raise RuntimeError( + f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) is not equal to world_size ({world_size})" + ) + # TODO (remi-or) the check above should probably happen during actual model sharding, same as the check below + # elif self.tp_size > 1 or self.fsdp_size > 1: + # issue = "not initialized" if torch.distributed.is_available() else "not available" + # raise ValueError(f"Requested {self.tp_size = }, {self.fsdp_size = } but torch.distributed is {issue}.") @classmethod - def from_dict(cls, config_dict, **kwargs): - """ - Constructs a DistributedConfig instance from a dictionary of parameters. - Args: - config_dict (Dict[str, Any]): Dictionary containing configuration parameters. - **kwargs: Additional keyword arguments to override dictionary values. - Returns: - DistributedConfig: Instance of DistributedConfig constructed from the dictionary. - """ - config = cls(**config_dict) - to_remove = [] - for key, value in kwargs.items(): - if hasattr(config, key): - setattr(config, key, value) - to_remove.append(key) - for key in to_remove: - kwargs.pop(key, None) - return config - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file + def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig": + merged = {**config_dict, **kwargs} + valid_keys = {f.name for f in cls.__dataclass_fields__.values()} + return cls(**{k: v for k, v in merged.items() if k in valid_keys}) + + def to_dict(self) -> dict: + return asdict(self) + + def to_json_string(self) -> str: + return json.dumps(self.to_dict(), indent=2) + "\n" + def to_json_file(self, json_file_path: str | os.PathLike): - """ - Save this instance to a JSON file. - Args: - json_file_path (`str` or `os.PathLike`): - Path to the JSON file in which this configuration instance's parameters will be saved. - use_diff (`bool`, *optional*, defaults to `True`): - If set to `True`, only the difference between the config instance and the default - `QuantizationConfig()` is serialized to JSON file. - """ - with open(json_file_path, "w", encoding="utf-8") as writer: - config_dict = self.to_dict() - json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - - writer.write(json_string) - - def to_dict(self) -> dict[str, Any]: - """ - Serializes this instance to a Python dictionary. Returns: - `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. - """ - return copy.deepcopy(self.__dict__) - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__ - def __iter__(self): - """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" - yield from copy.deepcopy(self.__dict__).items() - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ + with open(json_file_path, "w", encoding="utf-8") as f: + f.write(self.to_json_string()) + def __repr__(self): return f"{self.__class__.__name__} {self.to_json_string()}" - - def to_json_string(self): - """ - Serializes this instance to a JSON formatted string. - Returns: - str: JSON formatted string representing the configuration instance. - """ - return json.dumps(self.__dict__, indent=2) + "\n" - - def update(self, **kwargs): - """ - Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, - returning all the unused kwargs. - Args: - kwargs (`Dict[str, Any]`): - Dictionary of attributes to tentatively update this class. - Returns: - `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. - """ - to_remove = [] - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - to_remove.append(key) - - # Remove all the attributes that were updated, without modifying the input dict - unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} - return unused_kwargs diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py new file mode 100644 index 000000000000..64395ff2f531 --- /dev/null +++ b/src/transformers/distributed/fsdp.py @@ -0,0 +1,398 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import inspect +import os +from typing import TYPE_CHECKING, Any, Literal + +from ..utils import is_torch_available, is_torch_greater_or_equal, logging, strtobool +from ..utils.quantization_config import QuantizationMethod +from .tensor_parallel import replace_layer_number_by_wildcard + + +if TYPE_CHECKING: + import torch.nn as nn + +if is_torch_available(): + import torch + +if is_torch_available() and is_torch_greater_or_equal("2.6"): + from torch.distributed._composable.fsdp import fully_shard + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy + +logger = logging.get_logger(__name__) + + +def is_fsdp_enabled() -> bool: + """Check if FSDP is active via Accelerate (env var based) — covers FSDP1 only.""" + if not is_torch_available(): + return False + + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 + and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 + ) + + +def is_fsdp_managed_module(module: nn.Module) -> bool: + """Check if a module is managed by FSDP (1 or 2).""" + if not is_torch_available(): + return False + if not torch.distributed.is_available(): + return False + + # FSDP2: attribute set by apply_fsdp2() + if getattr(module, "_is_fsdp_managed_module", False): + return True + # FSDP1: wrapped by FullyShardedDataParallel + try: + from torch.distributed.fsdp import FullyShardedDataParallel + except ImportError: + return False + return isinstance(module, FullyShardedDataParallel) + + +def _get_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: + """Parse `cpu_offload` / `mixed_precision` flags from the user fsdp_plan into fully_shard kwargs.""" + policy_kwargs = {} + if fsdp_plan.get("cpu_offload"): + policy_kwargs["offload_policy"] = CPUOffloadPolicy() + if fsdp_plan.get("mixed_precision"): + policy_kwargs["mp_policy"] = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) + return policy_kwargs + + +def _parse_manual_plan_entry( + entry: list[str], +) -> tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: + """ + Returns: + tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: + - bool: whether to reshard after forward + - MixedPrecisionPolicy | None: mixed precision policy + - CPUOffloadPolicy | None: cpu offload policy + """ + + if not isinstance(entry, list): + raise ValueError( + f"Manual fsdp_plan values must be a list of strings combining strategy/policies, got {type(entry)}" + ) + items = entry + + strategy: Literal["free_full_weight", "keep_full_weight"] | None = None + offload_policy: CPUOffloadPolicy | None = None + mp_policy: MixedPrecisionPolicy | None = None + + for item in items: + if not isinstance(item, str): + raise ValueError( + f"fsdp_plan option must be a string, got {type(item)}. " + "Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." + ) + token = item.lower() + if token in {"free_full_weight", "keep_full_weight"}: + strategy = token + elif token == "cpu_offload": + offload_policy = CPUOffloadPolicy() + elif token == "mixed_precision": + # TODO(3outeille): add support for different dtypes + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=torch.bfloat16, + ) + else: + raise ValueError( + "Unknown fsdp_plan option " + f"{item!r}. Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." + ) + + if strategy is None: + strategy = "free_full_weight" + + return strategy != "keep_full_weight", mp_policy, offload_policy + + +def _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): + if pattern in name_to_module: + target = name_to_module[pattern] + if isinstance(target, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential)): + # (ModuleList, ModuleDict, Sequential) don't have a forward() that gets called - + # the model loops over their children directly. So when a pattern matches a + # container, we shard each child instead. + for child_name, child in target.named_children(): + yield f"{pattern}.{child_name}", child + else: + yield pattern, target + return + + # Prefix match: "model.layers" matches "model.layers.0", etc. + for name, module in model.named_modules(): + if name in already_sharded_names or isinstance( + module, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential) + ): + continue + if name != pattern and not name.startswith(pattern + "."): + continue + if any( + name.startswith(already_sharded_names_name + ".") for already_sharded_names_name in already_sharded_names + ): + continue + yield name, module + + +def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: + modules = fsdp_plan.get("modules") + if not isinstance(modules, dict): + raise ValueError("Manual fsdp_plan must define a 'modules' dict.") + return modules + + +def is_tail_pair(entries) -> bool: + """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" + if len(entries) != 2: + return False + names = [name for name, _ in entries] + has_norm = any(n == "norm" or n.endswith(".norm") for n in names) + has_head = any(n in {"lm_head", "embed_tokens"} or n.endswith((".lm_head", ".embed_tokens")) for n in names) + return has_norm and has_head + + +def tied_source_path(model) -> str | None: + """Return the dotted path of the input embedding module (the tied source).""" + input_embed = getattr(model, "get_input_embeddings", lambda: None)() + if input_embed is None: + return None + for name, module in model.named_modules(): + if module is input_embed: + return name + return None + + +def _resolve_plan_key(name_to_module: dict, key: str): + """Resolve a plan key into the matching (name, module) pairs. + + Supports exact module names and tp_plan-style wildcards (via + ``replace_layer_number_by_wildcard``). + """ + if key in name_to_module: + return [(key, name_to_module[key])] + return [(name, mod) for name, mod in name_to_module.items() if replace_layer_number_by_wildcard(name) == key] + + +def _iter_plan_targets(model, plan, is_weights_tied: bool, tied_source: str | None): + """Yield ``(name, module, strategy)`` for every module the plan applies to. + + Expands wildcards via ``_resolve_plan_key`` and pre-applies tying rules: + skips the standalone tied-source entry, and rewrites a keep ``"lm_head"`` + entry to the tied source so the shared parameter is wrapped once. + """ + name_to_module = dict(model.named_modules()) + for key, strategy in plan.items(): + if is_weights_tied and key == tied_source: + continue + if is_weights_tied and key == "lm_head" and strategy == "keep_full_weight" and tied_source is not None: + yield tied_source, name_to_module[tied_source], strategy + continue + for name, module in _resolve_plan_key(name_to_module, key): + yield name, module, strategy + + +def apply_fully_shard_data_parallel( + model, + fsdp_mesh, + fsdp_plan: dict[str, Any] | None, +): + """ + Apply FSDP2 (fully_shard) to a model. + + When ``fsdp_plan`` is ``None`` or doesn't contain a ``"modules"`` key, the + model-declared ``model._fsdp_plan`` drives sharding. Policies (`cpu_offload`, + `mixed_precision`) from ``fsdp_plan`` are applied on top. + + When ``fsdp_plan`` has a ``"modules"`` key, the user fully specifies the + layout (manual mode). + + Examples: + # Plan-driven (uses model._fsdp_plan). + fsdp_plan = None + fsdp_plan = {"cpu_offload": True, "mixed_precision": True} + + # Manual override. + fsdp_plan = { + "modules": { + "model.embed_tokens": ["free_full_weight"], + "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], + "model.norm": ["keep_full_weight"], + "lm_head": ["keep_full_weight"], + }, + } + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for FSDP support") + + if not is_torch_greater_or_equal("2.6"): + raise OSError("FSDP2 requires torch>=2.6") + + if fsdp_plan is None: + fsdp_plan = {} + + input_embed = getattr(model, "get_input_embeddings", lambda: None)() + output_embed = getattr(model, "get_output_embeddings", lambda: None)() + is_weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight + ) + + if not isinstance(fsdp_plan, dict): + raise ValueError(f"fsdp_plan must be a dict, got {type(fsdp_plan)}") + is_manual = "modules" in fsdp_plan + + if not is_manual: + policy_kwargs = _get_policy_kwargs(fsdp_plan) + + plan = getattr(model, "_fsdp_plan", None) or {} + if not plan: + raise ValueError( + f"{type(model).__name__} has no `_fsdp_plan` declared. Either set " + "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class, " + "or pass an explicit `fsdp_plan={'modules': {...}}` manual override." + ) + + tied_source = tied_source_path(model) if is_weights_tied else None + keep_buffer: list[tuple[str, Any]] = [] + + for name, module, strategy in _iter_plan_targets(model, plan, is_weights_tied, tied_source): + if strategy == "keep_full_weight": + keep_buffer.append((name, module)) + continue + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **policy_kwargs) + logger.debug(f"Applied fully_shard to {name} (reshard=True)") + + # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) + # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. + if is_tail_pair(keep_buffer): + keep_names = [n for n, _ in keep_buffer] + keep_modules = [m for _, m in keep_buffer] + fully_shard(keep_modules, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) + logger.debug(f"Grouped tail {keep_names} (reshard=False)") + else: + for name, module in keep_buffer: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) + logger.debug(f"Applied fully_shard to {name} (reshard=False)") + + # Shard root model + fully_shard(model, mesh=fsdp_mesh, **policy_kwargs) + + logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(plan)} entries") + + else: + # fsdp_plan = { + # "modules": { + # "model.layers.0.self_attn": ["free_full_weight"], # reshard_after_forward=True + # "model.norm": ["keep_full_weight"], # reshard_after_forward=False + # "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], + # }, + # } + + name_to_module = dict(model.named_modules()) + already_sharded_names: set[str] = set() + root_mp_policy = MixedPrecisionPolicy() + root_offload_policy = OffloadPolicy() + + for pattern, entry in _get_manual_plan_modules(fsdp_plan).items(): + reshard, mp_policy, offload_policy = _parse_manual_plan_entry(entry) + if mp_policy is not None: + root_mp_policy = mp_policy + if offload_policy is not None: + root_offload_policy = offload_policy + + for name, module in _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): + if name in already_sharded_names: + continue + shard_kwargs = {"mesh": fsdp_mesh, "reshard_after_forward": reshard} + if mp_policy is not None: + shard_kwargs["mp_policy"] = mp_policy + if offload_policy is not None: + shard_kwargs["offload_policy"] = offload_policy + fully_shard(module, **shard_kwargs) + already_sharded_names.add(name) + logger.debug(f"Applied fully_shard to {name}") + + # Shard root model with the same policies as sub-modules. + # MixedPrecisionPolicy.output_dtype casting happens in post_forward + # for every fully_shard-wrapped module, even with no direct parameters. + fully_shard(model, mesh=fsdp_mesh, mp_policy=root_mp_policy, offload_policy=root_offload_policy) + + # Used by generation code to detect FSDP and enable synced_gpus. + model._is_fsdp_managed_module = True + + if is_weights_tied and hasattr(model, "tie_weights"): + # Re-tie weights. + # fully_shard replaces nn.Parameter objects (swapping data for DTensor shards), + # which breaks weight tying (e.g. lm_head.weight is no longer embed_tokens.weight). + # Re-tying makes lm_head._parameters["weight"] point to the new DTensor parameter + # so gradients accumulate correctly into a single buffer. + model.tie_weights() + + return model + + +# ========================= PEFT compatibility ========================= +# TODO(3outeille): make sure new FSDP works with PEFT +def get_fsdp_ckpt_kwargs(): + """ + Returns checkpoint kwargs for FSDP model saving. + + Checks if the `adapter_only` parameter is supported by `save_fsdp_model` from accelerate + and returns the appropriate kwargs. + """ + from accelerate.utils import save_fsdp_model + + if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): + return {"adapter_only": True} + else: + return {} + + +def update_fsdp_plugin_peft(model, accelerator): + """ + Updates the FSDP plugin for PEFT LoRA/QLoRA compatibility. + + When using FSDP with PEFT LoRA, the auto wrap policy needs to be updated to additionally wrap + LoRA trainable layers separately. When using FSDP with QLoRA, the mixed precision policy needs + to be updated to use the quantization storage data type. + """ + from peft import PeftConfig + from peft.utils.other import fsdp_auto_wrap_policy + + if isinstance(model.active_peft_config, PeftConfig): + accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) + if ( + getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES + and model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point + ): + accelerator.state.fsdp_plugin.set_mixed_precision( + model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True + ) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py new file mode 100644 index 000000000000..03606a2768d5 --- /dev/null +++ b/src/transformers/distributed/sharding_utils.py @@ -0,0 +1,495 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +from ..utils import is_torch_available + + +if TYPE_CHECKING: + import torch + from torch.distributed.tensor import DTensor + +if is_torch_available(): + import torch + from torch.distributed._functional_collectives import wait_tensor + from torch.distributed.tensor import DTensor, Replicate + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + from torch.distributed.tensor.placement_types import Shard, _StridedShard + + # torch < 2.10 names as an underscore before `local_shard_size_and_offset`: alias it the non-underscored version + if not hasattr(Shard, "local_shard_size_and_offset") and hasattr(Shard, "_local_shard_size_and_offset"): + Shard.local_shard_size_and_offset = Shard._local_shard_size_and_offset + + +class DtensorShardOperation: + """Shard-on-read: slice a full checkpoint tensor down to this rank's local + DTensor shard, for any combination of placements on a 1-D or n-D mesh. + + Placements primer + ----------------- + Each mesh dim carries one placement describing how it slices the tensor: + + | Placement | Local data on each rank of the mesh dim | + |--------------------------|-------------------------------------------------| + | Replicate | full tensor (no slicing) | + | Shard(d) | contiguous chunk of dim d (rows r*c .. (r+1)*c) | + | _StridedShard(d, sf=N) | one chunk from each of N groups along dim d, | + | | concatenated together (interleaved layout) | + + Different scenarios of different placements + ------------------------------------------ + Placement tuples are ordered outermost-first; for a 2-D (fsdp, tp) mesh + the tuple is (fsdp_placement, tp_placement). + + | Scenario | Placements | + |---------------------------------------------------|---------------------------------------| + | TP-only, non-fused (e.g. q_proj/k_proj/v_proj) | [Shard(d)] | + | TP-only, fused gate/up | [_StridedShard(d, sf=2)] | + | TP + FSDP, same tensor dim | [_StridedShard(d, sf=tp_size), Shard(d)] | + | TP + FSDP, different dims | [Shard(d1), Shard(d2)] | + + Mesh dimensions are listed outermost-first. For a 2-D (fsdp=F, tp=T) mesh, + rank index = fsdp_idx * T + tp_idx. + + Loading (this class) + -------------------- + During `from_pretrained`, each rank reads the full checkpoint tensors, + then calls `shard_tensor(source, tensor_idx=...)` to keep only its local + DTensor shard. The class encapsulates the placements + mesh so the + slicing logic doesn't have to be repeated at every call site. + + Two checkpoint layouts are supported. Running example: an MoE weight + stack of param shape [N_experts, in, out]: + + | Checkpoint layout | tensor_idx | source.shape | Returns | + |------------------------------------|------------|---------------|-----------------------------| + | One stacked tensor | None | [N, in, out] | This rank's slice along | + | | | | every sharded dim. | + | N per-expert tensors (called once | 0..N-1 | [in, out] | Inner-dim slice if this | + | per expert by the caller) | | | rank owns expert `i`, | + | | | | else None (caller drops it).| + """ + + def __init__(self, param: DTensor): + self.device_mesh = param.device_mesh + self.placements = tuple(param.placements) + self.param_ndim = param.ndim + local_shape, offsets = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) + # Where this rank's slice starts along axis 0, and how many indices + # it covers. Example: param of shape [8, in, out] with Shard(0) on + # 2 ranks gives: + # rank 0 → _axis0_offset=0, _axis0_local_size=4 (owns experts 0..3) + # rank 1 → _axis0_offset=4, _axis0_local_size=4 (owns experts 4..7) + # When the checkpoint stores one tensor per expert, shard_tensor + # checks whether tensor_idx falls in this rank's range to decide + # whether to keep the piece or drop it. + self._axis0_offset = offsets[0] + self._axis0_local_size = local_shape[0] + + def shard_tensor( + self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None + ) -> torch.Tensor | None: + """Slice source down to this rank's shard. + + The checkpoint can store the parameter in two layouts. Take a stack + of N MoE experts of shape [in, out] as a running example — the + param shape is [N, in, out]: + + - Single tensor (tensor_idx is None): the checkpoint holds one + [N, in, out] tensor, so source.shape == param.shape. Every + sharded dim is sliced here, including axis 0. + + - One tensor per piece (tensor_idx given): the checkpoint holds N + separate [in, out] tensors, one per expert. shard_tensor is + called once per expert; on each call source is the [in, out] + tensor for expert number tensor_idx (so 0 <= tensor_idx < N). + Note: source has one fewer dim than the param: the axis-0 + index lives in tensor_idx, not in source.shape. + If this rank does not own tensor_idx along axis 0, return None + and the piece is discarded. Otherwise slice only the inner + dims; the caller (MergeModulelist / Concatenate) collects the + kept pieces and stacks them back along axis 0 to rebuild the + full [N, in, out] param. + """ + source_shape = list(source.shape) if isinstance(source, torch.Tensor) else source.get_shape() + placements = [(md, p) for md, p in enumerate(self.placements) if hasattr(p, "dim")] + + # Dense path + if tensor_idx is None: + if not placements: + return source[...].to(device=device, dtype=dtype) + has_strided = any(not p.is_shard() for _, p in placements) + intervals = [[(0, size)] for size in source_shape] + for mesh_dim, placement in placements: # [i.e: (0, Shard(0)), (1, Shard(-1))] + sub_mesh = self._get_sub_mesh(mesh_dim) + rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() + source_dim = self._norm_dim(placement.dim) + if not placement.is_shard(): + intervals[source_dim] = self._strided_intervals( + intervals[source_dim], rank, world_size, placement.split_factor + ) + else: + intervals[source_dim] = self._contiguous_intervals(intervals[source_dim], rank, world_size) + # Only _StridedShard can produce multi-interval dims that need cat. + if has_strided: + return self._slice_and_cat(source, intervals, device, dtype) + slices = tuple(slice(*(pieces[0] if pieces else (0, 0))) for pieces in intervals) + return source[slices].to(device=device, dtype=dtype) + + # MoE path: drop the piece if this rank does not own tensor_idx + # along axis 0. Once shard_tensor has been called for all N pieces, + # the caller (MergeModulelist) stacks the kept slices along axis 0 to + # form this rank's local shard of the param. + shards_leading_axis = any(self._norm_dim(p.dim) == 0 for _, p in placements) + owns_index = self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size + if shards_leading_axis and not owns_index: + return None + + # Inner dims use only _contiguous_intervals (one piece per dim), so a + # single slice suffices + inner_placements = [(md, p) for md, p in placements if self._norm_dim(p.dim) != 0] + if not inner_placements: + return source[...].to(device=device, dtype=dtype) + slice_per_dim: list[tuple[int, int]] = [(0, size) for size in source_shape] + # placement.dim is indexed in param space (e.g. axis 2 of [N, in, out]). + # source is in source space (e.g. axis 1 of [in, out]), so we translate + # from one to the other by stripping the leading axis. + for mesh_dim, placement in inner_placements: + sub_mesh = self._get_sub_mesh(mesh_dim) + rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() + param_dim = self._norm_dim(placement.dim) + source_dim = param_dim - 1 + pieces = self._contiguous_intervals([slice_per_dim[source_dim]], rank, world_size) + slice_per_dim[source_dim] = pieces[0] if pieces else (0, 0) + return source[tuple(slice(s, e) for s, e in slice_per_dim)].to(device=device, dtype=dtype) + + def _strided_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int, split_factor: int + ) -> list[tuple[int, int]]: + narrowed = [] + for start, end in intervals: + group_size = math.ceil((end - start) / split_factor) + for group_idx in range(split_factor): + group_start = start + group_idx * group_size + group_end = min(group_start + group_size, end) + if group_end <= group_start: + continue + size, offset = Shard.local_shard_size_and_offset(group_end - group_start, world_size, rank) + if size > 0: + narrowed.append((group_start + offset, group_start + offset + size)) + return narrowed + + def _contiguous_intervals( + self, intervals: list[tuple[int, int]], rank: int, world_size: int + ) -> list[tuple[int, int]]: + total = sum(end - start for start, end in intervals) + my_size, my_offset = Shard.local_shard_size_and_offset(total, world_size, rank) + if my_size == 0: + return [] + + out: list[tuple[int, int]] = [] + flat_pos = 0 + slice_end = my_offset + my_size + for start, end in intervals: + length = end - start + interval_end_flat = flat_pos + length + if interval_end_flat <= my_offset: # entirely before my slice + flat_pos = interval_end_flat + continue + if flat_pos >= slice_end: # entirely after my slice + break + sub_start = max(0, my_offset - flat_pos) + sub_end = min(length, slice_end - flat_pos) + out.append((start + sub_start, start + sub_end)) + flat_pos = interval_end_flat + return out + + def _slice_and_cat(self, source, intervals, device, dtype): + multi_interval_dim: int | None = None + slices: list[slice] = [] + for source_dim, pieces in enumerate(intervals): + if len(pieces) == 1: + start, end = pieces[0] + slices.append(slice(start, end)) + continue + if multi_interval_dim is not None: + raise ValueError("Shard-on-read only supports disjoint ranges on a single checkpoint dimension.") + multi_interval_dim = source_dim + slices.append(slice(None)) # placeholder, filled per-piece below + + # Fast path: every dim is one contiguous interval, read in a single slice. + if multi_interval_dim is None: + return source[tuple(slices)].to(device=device, dtype=dtype) + + # Multi-interval dim: read each piece separately, then concatenate. + pieces_read = [] + for start, end in intervals[multi_interval_dim]: + piece_slices = list(slices) + piece_slices[multi_interval_dim] = slice(start, end) + pieces_read.append(source[tuple(piece_slices)]) + return torch.cat(pieces_read, dim=multi_interval_dim).to(device=device, dtype=dtype) + + def _get_sub_mesh(self, mesh_dim: int): + if self.device_mesh.ndim == 1: + return self.device_mesh + return self.device_mesh[self.device_mesh.mesh_dim_names[mesh_dim]] + + def _norm_dim(self, dim: int) -> int: + # if dim is negative, it should be normalized to the last axis + return dim if dim >= 0 else self.param_ndim + dim + + +def _dtensor_from_local_like(local_tensor: torch.Tensor, ref: DTensor) -> DTensor: + """Wrap `local_tensor` as a DTensor that mirrors `ref`'s mesh, placements, + global shape, and stride.""" + return DTensor.from_local( + local_tensor.contiguous(), + ref.device_mesh, + ref.placements, + run_check=False, + shape=ref.shape, + stride=tuple(ref.stride()), + ) + + +def _replicate_dtensor(tensor: DTensor) -> DTensor: + """All-gather a DTensor to fully Replicate, handling _StridedShard. + + PyTorch's redistribute() does not support _StridedShard as a source: + _StridedShard -> redistribute() -> Replicate ❌ AssertionError + _StridedShard -> redistribute() -> Shard ❌ NotImplementedError + Shard -> redistribute() -> Replicate ✅ works + Replicate -> redistribute() -> Shard ✅ works + Replicate -> redistribute() -> _StridedShard ✅ works + + During reconstruction, we walk placements right-to-left (innermost mesh dim first), + invoke each one's low-level _to_replicate_tensor, and wait for the async collective to finish + at each step. Example — global [256, 1024], mesh (fsdp=2, tp=2), + placements (_StridedShard(0), Shard(0)), local [64, 1024]: + + i=1 (tp, Shard(0)): [64, 1024] -> [128, 1024] + i=0 (fsdp, _StridedShard(0)): [128, 1024] -> [256, 1024] + """ + mesh = tensor.device_mesh + placements = tensor.placements + replicate_all = tuple(Replicate() for _ in range(mesh.ndim)) + + if not any(isinstance(p, _StridedShard) for p in placements): + return tensor.redistribute(placements=replicate_all) + + with torch.no_grad(): + local = tensor._local_tensor + for i in reversed(range(mesh.ndim)): + p = placements[i] + if p.is_replicate(): + continue + logical_shape = list(tensor.shape) + for j, pj in enumerate(placements[:i]): + if not pj.is_replicate(): + size, _ = Shard.local_shard_size_and_offset( + logical_shape[pj.dim], mesh.size(j), mesh.get_local_rank(j) + ) + logical_shape[pj.dim] = size + local = p._to_replicate_tensor(local, mesh, i, logical_shape) + local = wait_tensor(local) + return DTensor.from_local(local, mesh, replicate_all, run_check=False) + + +def _find_strided_shard_placement_from_fused_params(placements): + for i, p in enumerate(placements): + if not isinstance(p, _StridedShard): + continue + # We want to find the first _StridedShard placement that is not composed with another placement on the same dim. + # Because that means it's a fused parameter. Meanwhile if it's acting on the same dim, that means it comes from composing parallelism together + has_partner_on_same_dim = any( + j != i and getattr(other, "dim", None) == p.dim for j, other in enumerate(placements) + ) + if not has_partner_on_same_dim: + return p + return None + + +def _split_fused_dtensor(dt, dim, n_pieces): + local_pieces = list(dt._local_tensor.chunk(n_pieces, dim=dim)) + if len(local_pieces) != n_pieces: + raise RuntimeError( + f"Cannot split DTensor of shape {tuple(dt.shape)} into {n_pieces} pieces along dim {dim}: " + f"got {len(local_pieces)} chunks instead." + ) + new_placements = tuple( + Shard(p.dim) if (isinstance(p, _StridedShard) and p.dim == dim) else p for p in dt.placements + ) + return [ + DTensor.from_local(lp.contiguous(), dt.device_mesh, new_placements, run_check=False) for lp in local_pieces + ] + + +def _merge_unfused_dtensors(pieces, dim, target_placements): + local_cat = torch.cat([p._local_tensor for p in pieces], dim=dim).contiguous() + return DTensor.from_local(local_cat, pieces[0].device_mesh, target_placements, run_check=False) + + +def get_fusion_metadata(optimizer_state_dict): + """Inspect the optimizer state dict and return metadata describing every + fused param that needs to be split for DCP. Does NOT mutate + `optimizer_state_dict`. Pair it with `unfuse_optimizer_state` to apply + the split and `fuse_optimizer_state` to undo it on load. + + Args: + optimizer_state_dict: as returned by `get_optimizer_state_dict(...)`. For + a Mixtral on a (fsdp=2, tp=2) mesh with AdamW and a single MoE layer, + it looks like: + + { + "state": { + "model.layers.0.mlp.experts.gate_up_proj": { + "exp_avg": DTensor(shape=(4, 16, 8), + placements=(Shard(0), _StridedShard(1, sf=2))), + "exp_avg_sq": DTensor(shape=(4, 16, 8), + placements=(Shard(0), _StridedShard(1, sf=2))), + "step": tensor(7.0), + }, + "model.layers.0.mlp.experts.down_proj": { + "exp_avg": DTensor(shape=(4, 8, 16), + placements=(Shard(0), Shard(2))), + "step": tensor(7.0), + }, + }, + "param_groups": [{"lr": 1e-4, ...}], + } + + Returns: + fusion_metadata = { + "model.layers.0.mlp.experts.gate_up_proj": { + "chunk_dim": 1, + "placements": (Shard(0), _StridedShard(1, split_factor=2)), + "unfused_keys": [ + "model.layers.0.mlp.experts.gate_up_proj.0", + "model.layers.0.mlp.experts.gate_up_proj.1", + ], + }, + } + """ + optimizer_state = optimizer_state_dict.get("state", {}) + fusion_metadata: dict[str, dict] = {} + + for param_name, optimizer_fields in optimizer_state.items(): + dtensor = next((v for v in optimizer_fields.values() if isinstance(v, DTensor)), None) + if dtensor is None: + continue + strided_shard_placement = _find_strided_shard_placement_from_fused_params(dtensor.placements) + if strided_shard_placement is None: + continue + fusion_metadata[param_name] = { + "chunk_dim": strided_shard_placement.dim, + "placements": tuple(dtensor.placements), + "unfused_keys": [f"{param_name}.{i}" for i in range(strided_shard_placement.split_factor)], + } + + return fusion_metadata + + +def unfuse_optimizer_state(optimizer_state_dict, fusion_metadata): + """Apply `fusion_metadata` to split each fused param into its `sf` plain-Shard + pieces in place. After the call, `optimizer_state_dict["state"]` has + `.0` .. `.{sf-1}` keys in place of each original fused key; + non-fused params (absent from `fusion_metadata`) are untouched. + + Args: + optimizer_state_dict: same shape as the input to `get_fusion_metadata`. + fusion_metadata: as returned by `get_fusion_metadata(optimizer_state_dict)`. + """ + optimizer_state = optimizer_state_dict.get("state", {}) + + for param_name, metadata in fusion_metadata.items(): + optimizer_fields = optimizer_state[param_name] + + for unfused_key in metadata["unfused_keys"]: + optimizer_state[unfused_key] = {} + + for optim_param_name, value in optimizer_fields.items(): + if isinstance(value, DTensor): + chunks = _split_fused_dtensor(value, metadata["chunk_dim"], len(metadata["unfused_keys"])) + else: + chunks = [value] * len(metadata["unfused_keys"]) # scalar — replicate + for unfused_key, chunk in zip(metadata["unfused_keys"], chunks): + optimizer_state[unfused_key][optim_param_name] = chunk + + del optimizer_state[param_name] + + +def fuse_optimizer_state(optimizer_state_dict, fusion_metadata): + """Fuse the optimizer state dict back in place using the fusion info. + + Inverse of `unfuse_optimizer_state`: concatenates each set of per-piece + sub-dicts along `chunk_dim`, rewraps each DTensor field with the original + `_StridedShard` placement, then replaces the piece keys with the fused key + in `optimizer_state_dict["state"]`. + + Args: + optimizer_state_dict: the unfused state dict (e.g. just filled by + `dcp.load`). For our Mixtral example, before this call it looks like: + + { + "state": { + "model.layers.0.mlp.experts.gate_up_proj.0": { + "exp_avg": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "exp_avg_sq": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "step": tensor(7.0), + }, + "model.layers.0.mlp.experts.gate_up_proj.1": { + "exp_avg": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "exp_avg_sq": DTensor(shape=(4, 8, 8), + placements=(Shard(0), Shard(1))), + "step": tensor(7.0), + }, + "model.layers.0.mlp.experts.down_proj": { + "exp_avg": DTensor(shape=(4, 8, 16), + placements=(Shard(0), Shard(2))), + "step": tensor(7.0), + }, + }, + "param_groups": [{"lr": 1e-4, ...}], + } + + After the call the `.0`/`.1` piece keys are gone and the original + fused key is restored with the `_StridedShard` placement. + + fusion_metadata: as returned by `get_fusion_metadata`. + """ + optimizer_state = optimizer_state_dict.get("state", {}) + + for param_name, metadata in fusion_metadata.items(): + first_piece = optimizer_state[metadata["unfused_keys"][0]] + + merged_fields = {} + for optim_param_name, value in first_piece.items(): + chunks = [optimizer_state[unfused_key][optim_param_name] for unfused_key in metadata["unfused_keys"]] + if isinstance(value, DTensor): + merged_fields[optim_param_name] = _merge_unfused_dtensors( + chunks, metadata["chunk_dim"], metadata["placements"] + ) + else: + merged_fields[optim_param_name] = value # scalar — pick any copy + + optimizer_state[param_name] = merged_fields + + for unfused_key in metadata["unfused_keys"]: + del optimizer_state[unfused_key] diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py new file mode 100644 index 000000000000..126b404594f2 --- /dev/null +++ b/src/transformers/distributed/tensor_parallel.py @@ -0,0 +1,526 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import contextlib +import re + +from ..utils import logging +from ..utils.generic import GeneralInterface +from ..utils.import_utils import is_torch_available, is_torch_greater_or_equal + + +if is_torch_available(): + import torch + +if is_torch_available() and is_torch_greater_or_equal("2.5"): + import torch.distributed as dist + from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor + from torch.distributed.tensor.parallel import ( + ColwiseParallel, + PrepareModuleInput, + RowwiseParallel, + SequenceParallel, + ) + from torch.distributed.tensor.parallel.style import ParallelStyle + from torch.distributed.tensor.placement_types import _StridedShard + + # Cache this result has it's a C FFI call which can be pretty time-consuming + _torch_distributed_available = torch.distributed.is_available() + + +logger = logging.get_logger(__name__) + + +def replace_layer_number_by_wildcard(name: str) -> str: + """ + Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between + a dot (`.`) and the end of the string. + This matches how modules are named/numbered when using a nn.ModuleList or nn.Sequential, but will NOT match + numbers in a parameter name itself, e.g. if the param is named `"w1"` or `"w2"`. + """ + return re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), name) + + +def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None: + """ + Get the TP style for a parameter from the TP plan. + + The TP plan is a dictionary that maps parameter names to TP styles. + The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight"). + + The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but + not parent classes for `post_init` calls + """ + generic_param_name = replace_layer_number_by_wildcard(parameter_name) + if generic_param_name in tp_plan: + return tp_plan[generic_param_name] + elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + return None + + +# ============================================================================= +# High-Level API Functions +# ============================================================================= + + +def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): + """ + Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. + + Only weight-sharding rules (colwise, rowwise, vocab, moe_experts) are checked. + Module/activation entries (e.g. PrepareModuleInput, SequenceParallel) set up + communication hooks on modules, not weight sharding, so they are excluded. + """ + + if tp_plan is None: + return + + # Filter out module-level comm hooks — they don't shard weights. + # Plan values are registry names; entries beginning with "activation" or "module" + # configure communication hooks rather than parameter sharding. + weight_plan = { + k: v for k, v in tp_plan.items() if not (v == "activation" or v.startswith(("activation_", "module_"))) + } + + generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} + unsharded_layers = set(generic_keys) + unused_rules = weight_plan.copy() + + for key in generic_keys: + param_name = key.rsplit(".", 1)[0] if "." in key else key + generic_param_name = re.sub(r"\d+", "*", param_name) + + if generic_param_name in weight_plan: + unused_rules.pop(generic_param_name, None) + unsharded_layers.discard(key) + elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in weight_plan: + unused_rules.pop(parent_param_name, None) + unsharded_layers.discard(key) + + if len(unused_rules) > 0: + logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}") + if len(unsharded_layers) > 0: + logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") + + +class TensorParallelStyle(ParallelStyle): + """Base class for transformers TP styles. Installs the pre / around / post + forward hooks. Subclasses that need to shard params override `_apply` to + wrap them as DTensor placeholders before calling `super()._apply(...)`. + + Param wrapping runs on meta (the model is on meta when `apply_tensor_parallel` + is invoked); `distribute_tensor` on meta builds metadata only — no collective. + Real data flows in later, async, via DtensorShardOperation during load. + + Forward-time hooks (override what you need): + - transform_inputs_pre_forward(module, args, kwargs, mesh) → (args, kwargs) + - context_around_forward(module) → context manager wrapping the call + - transform_output_post_forward(module, output, mesh) → output + """ + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + return args, kwargs + + def context_around_forward(self, module): + return contextlib.nullcontext() + + def transform_output_post_forward(self, module, output, mesh): + return output + + def _apply(self, module, mesh): + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh) + with self.context_around_forward(module): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh) + + module.forward = tp_forward + return module + + +class PrepareModuleInputOutput(TensorParallelStyle): + """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). + + Used for MoE blocks with SP: the input sequence is gathered before routing, + and the output (after expert allreduce) is split back to match the residual. + Forward output split is a local op (no comm). Backward creates the all-gather. + """ + + def __init__(self, use_local_output=True): + super().__init__() + self.use_local_output = use_local_output + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [Shard(1)], run_check=False) + x = x.redistribute(placements=[Replicate()]).to_local() + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [Replicate()], run_check=False) + return output.redistribute(placements=[Shard(1)]).to_local() + + +def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: + """Stitch a local grad back onto the original DTensor parameter. + + During forward we replace the DTensor param with a detached plain-tensor + leaf (see ``_swap_dtensor_params_for_local``) because ``grouped_mm`` / fused + ops don't accept DTensor inputs. That swap breaks the autograd link between + the local leaf's grad and the DTensor param's ``.grad``, so this tensor hook + runs on the leaf and copies/accumulates the grad onto the original DTensor. + + NOTE: An autograd-aware ``param.to_local()`` swap would let backward stitch + the grad automatically, but DTensor's backward path then redistributes the + resulting grad — and that redistribute does not currently support + ``_StridedShard`` placements (used by ``MoEExpertsParallel`` / ``PackedColwiseParallel``). + """ + tensor_meta = original_param._spec.tensor_meta + detached_grad = local_grad.detach() + grad_dtensor = DTensor.from_local( + detached_grad, + original_param.device_mesh, + original_param.placements, + run_check=False, + shape=tensor_meta.shape, + stride=tensor_meta.stride, + ) + with torch.no_grad(): + existing_grad = original_param.grad + if existing_grad is None: + original_param.grad = grad_dtensor + elif isinstance(existing_grad, DTensor): + existing_grad._local_tensor.add_(detached_grad) + else: + existing_grad.add_(detached_grad) + return local_grad + + +@contextlib.contextmanager +def _swap_dtensor_params_for_local(module): + """Temporarily replace DTensor params with local-shard ``Parameter``s for forward. + + ``grouped_mm`` / fused kernels don't accept DTensor inputs, so each DTensor + param is swapped for a detached local ``Parameter``. A tensor hook + (``_accumulate_local_param_grad``) on the local leaf copies the backward + grad back onto the original DTensor. + + The original DTensor params are restored on exit (even on exception) so + save_pretrained / state-dict still see sharded params. + """ + shadows = {} + for name, param in list(module.named_parameters(recurse=False)): + if not isinstance(param, DTensor): + continue + shadows[name] = param + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) + module._parameters.pop(name) + setattr(module, name, local) + try: + yield + finally: + for name, param in shadows.items(): + if hasattr(module, name): + delattr(module, name) + module.register_parameter(name, param) + + +class PackedColwiseParallel(TensorParallelStyle): + """Column-wise parallel style for fused linear weights packed along the output dimension.""" + + def __init__( + self, + *, + input_layouts=None, + use_local_output: bool = True, + split_factor: int = 2, + ): + super().__init__() + self.input_layouts = (input_layouts or Replicate(),) + self.use_local_output = use_local_output + self.split_factor = split_factor + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + input_tensor = args[0] + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, mesh, self.input_layouts, run_check=False) + elif input_tensor.placements != self.input_layouts: + input_tensor = input_tensor.redistribute(placements=self.input_layouts) + input_tensor = input_tensor.to_local() + return (input_tensor,) + args[1:], kwargs + + def context_around_forward(self, module): + return _swap_dtensor_params_for_local(module) + + def transform_output_post_forward(self, module, output, mesh): + if output is None or self.use_local_output: + return output + return DTensor.from_local( + output, mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False + ) + + def _apply(self, module, mesh): + if not isinstance(module, torch.nn.Linear): + raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") + # Wrap weight + bias as DTensor placeholders. Runs on meta — + # distribute_tensor builds metadata only, no collective. + placement = _StridedShard(dim=0, split_factor=self.split_factor) + for name in ("weight", "bias"): + meta = module._parameters.get(name) + if meta is None: + continue + module._parameters[name] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + return super()._apply(module, mesh) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"use_local_output={self.use_local_output}, split_factor={self.split_factor})" + ) + + +if is_torch_available() and is_torch_greater_or_equal("2.5"): + + class _AllReduceBackward(torch.autograd.Function): + """Identity forward, allreduce-sum backward. + + Used for MoE routing weights: the forward value is replicated (same on all + ranks), but the backward gradient is partial (each rank has 1/tp_size from + its expert shard). We need to sum the partial gradients without dividing by + world_size, which is what DTensor's ``Replicate`` backward does incorrectly. + """ + + @staticmethod + def forward(ctx, x, process_group): + ctx.process_group = process_group + return x + + @staticmethod + def backward(ctx, grad): + dist.all_reduce(grad, group=ctx.process_group) + return grad, None + + +class MoEExpertsParallel(TensorParallelStyle): + """Tensor-parallel style for MoE expert modules. + + Shards expert weights as DTensors, then wraps the module's ``forward`` so + that grouped_mm (which needs plain tensors) works transparently. + + Lifecycle phases: + 1. _apply — wrap each expert weight named in shard_plan as a DTensor + placeholder with the declared placement. + 2. transform_inputs_pre_forward — localize hidden_states (Replicate→local, + gives us an all-reduce on the backward gradient for free), then fix + routing-weight gradients (their backward is partial; use allreduce-sum, + not divide-by-world-size). + 3. context_around_forward — swap DTensor params for local leaves so + grouped_mm sees plain tensors; restored on exit so save_pretrained + still sees DTensors. + 4. transform_output_post_forward — under TP-only each rank's output is + partial (only its expert shard contributed); reduce/redistribute to + output_layouts. + """ + + def __init__(self, output_layouts=None, shard_plan=None): + super().__init__() + self.output_layouts = output_layouts or Replicate() + self._moe_shard_plan = shard_plan or {} + + def _apply(self, module, mesh): + # Wrap each expert weight as a DTensor placeholder. Runs on meta — + # distribute_tensor builds metadata only, no collective. + for name, placement in self._moe_shard_plan.items(): + meta = module._parameters.get(name) + if meta is None: + continue + module._parameters[name] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + return super()._apply(module, mesh) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + hidden_states, top_k_index, top_k_weights = args + if not isinstance(hidden_states, DTensor): + hidden_states = DTensor.from_local(hidden_states, mesh, [Replicate()], run_check=False) + hidden_states = hidden_states.to_local() + + if isinstance(top_k_weights, DTensor): + top_k_weights = top_k_weights.to_local() + tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + + return (hidden_states, top_k_index, top_k_weights), kwargs + + def context_around_forward(self, module): + return _swap_dtensor_params_for_local(module) + + def transform_output_post_forward(self, module, output, mesh): + if output is None: + return None + # Under TP-only each rank has a partial result; under TP+FSDP the + # weights may be fully gathered by FSDP, making the output complete. + has_sharded_params = any( + isinstance(p, DTensor) and any(not pl.is_replicate() for pl in p.placements) for p in module.parameters() + ) + source = Partial() if has_sharded_params else Replicate() + if not isinstance(output, DTensor): + output = DTensor.from_local(output, mesh, [source], run_check=False) + # MoE output is 2D [tokens, hidden]. For SP, Shard(1) means seq dim + # in 3D but token dim (0) in 2D. + target = self.output_layouts + if output.dim() == 2 and isinstance(target, Shard) and target.dim == 1: + target = Shard(0) + if output.placements != (target,): + output = output.redistribute(placements=(target,)) + return output.to_local() + + +class ParallelInterface(GeneralInterface): + """Registry of named TP styles. Configs and modeling files reference these by string name. + + Adding a new entry here is the supported way to introduce a new TP style. + Users can also override or extend at runtime via ``ALL_PARALLEL_STYLES["my_style"] = ...``. + + Naming convention: ``{kind}[_{comm}][_{extra}]``. The ``_{comm}`` suffix is dropped only when + comm is ``"none"`` (no collective). All entries are eager instances; the dict literal lives + behind a torch-availability guard so this module remains importable without torch. + """ + + _global_mapping = ( + { + # Column-parallel + "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), + "colwise_allgather": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "colwise_loss_parallel": ColwiseParallel( + input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False + ), + "packed_colwise": PackedColwiseParallel(input_layouts=Replicate()), + # Row-parallel + "rowwise_allreduce": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), + "rowwise_reduce_scatter": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Shard(1)), + # Vocab / embedding (rowwise sharding on vocab dim) + "vocab_allreduce": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "vocab_reduce_scatter": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)), + # Activation / norm (sequence-parallel passthrough) + # use_local_output=True: torch defaults to False here, but downstream modeling + # code expects plain tensors, not DTensors. + "activation": SequenceParallel(use_local_output=True), + "activation_seq_dim_2": SequenceParallel(sequence_dim=2, use_local_output=True), + # Module-level prepare-input. Same use_local_output=True override as above — + # torch's default is False, our modeling code expects plain tensors downstream. + "module_allgather": PrepareModuleInput( + input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),), use_local_output=True + ), + "module_allgather_hidden_states": PrepareModuleInput( + input_kwarg_layouts={"hidden_states": Shard(1)}, + desired_input_kwarg_layouts={"hidden_states": Replicate()}, + use_local_output=True, + ), + "module_allgather_split": PrepareModuleInputOutput(), + # MoE — canonical shard_plan baked in (only variant in use across configs). + # gate_up_proj is packed (gate||up along output dim) so we use _StridedShard + # to interleave; down_proj is plain rowwise on its input dim. + "moe_experts_allreduce": MoEExpertsParallel( + output_layouts=Replicate(), + shard_plan={ + "gate_up_proj": _StridedShard(dim=-2, split_factor=2), + "down_proj": Shard(-1), + }, + ), + } + if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available + else {} + ) + + +ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() + + +def apply_tensor_parallel(model, tp_mesh, tp_plan): + """Apply tensor parallelism by calling each style's ``_apply`` on the + matching submodules. + + Walks ``model.named_modules()``, resolves each name against the wildcard + ``tp_plan`` from the model config, and applies the corresponding style + from ``ALL_PARALLEL_STYLES`` (looked up by string name) directly. + """ + distributed_config = getattr(model.config, "distributed_config", None) + sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) + sp_supported = getattr(model.config, "base_model_sp_plan", None) is not None + enable_sp = sp_requested and sp_supported + + if tp_plan is None: + if enable_sp: + tp_plan = dict(model._sp_plan or {}) + else: + tp_plan = dict(model._tp_plan or {}) + + # tie_weights() replaces lm_head.weight with embed_tokens.weight after TP is applied. + # If embed_tokens isn't in the plan, sharding lm_head as a DTensor causes tie to + # clobber it with a plain tensor (and forward then mixes DTensor/Tensor). Skip + # lm_head TP in that case so both ends stay plain and the tie is a real alias. + if getattr(model.config, "tie_word_embeddings", False): + tied_source_in_plan = any(k.endswith("embed_tokens") for k in tp_plan) + if not tied_source_in_plan: + tp_plan.pop("lm_head", None) + + for name, submodule in model.named_modules(): + style_value = _get_parameter_tp_plan(parameter_name=name, tp_plan=tp_plan, is_weight=False) + if style_value is None: + continue + ALL_PARALLEL_STYLES[style_value]._apply(submodule, tp_mesh) + + # Under SP, inputs_embeds is sequence-sharded after embed_tokens, so + # auto-generated position_ids would use the wrong (local) seq_len. + # Inject position_ids from the original input_ids shape before the model forward + if enable_sp: + base_model = getattr(model, model.base_model_prefix, model) + + def _inject_sp_metadata(mod, args, kwargs): + input_ids = kwargs.get("input_ids", args[0] if args else None) + if input_ids is None: + return args, kwargs + if "position_ids" not in kwargs or kwargs["position_ids"] is None: + seq_len = input_ids.shape[1] + kwargs["position_ids"] = torch.arange(seq_len, device=input_ids.device).unsqueeze(0) + return args, kwargs + + base_model.register_forward_pre_hook(_inject_sp_metadata, with_kwargs=True) + + # If the plan uses loss_parallel on lm_head, enable it globally so + # the model's internal loss computation handles DTensor logits correctly. + # loss_parallel patches F.cross_entropy to work with Shard(-1) logits. + # It must be active during both forward and backward, so we enable it + # once rather than as a context manager. + has_loss_parallel = any(v == "colwise_loss_parallel" for v in tp_plan.values()) + if has_loss_parallel: + from torch.distributed.tensor.parallel import loss_parallel + + model._loss_parallel_ctx = loss_parallel() + model._loss_parallel_ctx.__enter__() + + return model diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py new file mode 100644 index 000000000000..ae159e510a7d --- /dev/null +++ b/src/transformers/distributed/utils.py @@ -0,0 +1,299 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from ..utils import is_torch_available, is_torch_greater_or_equal, logging +from .fsdp import apply_fully_shard_data_parallel +from .sharding_utils import ( + _find_strided_shard_placement_from_fused_params, + _replicate_dtensor, + fuse_optimizer_state, + get_fusion_metadata, + unfuse_optimizer_state, +) +from .tensor_parallel import apply_tensor_parallel + + +logger = logging.get_logger(__name__) + + +if TYPE_CHECKING: + import torch.nn as nn + + from .configuration_utils import DistributedConfig + +if is_torch_available(): + import torch + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.state_dict import ( + get_model_state_dict, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + from torch.distributed.tensor import DTensor + + if is_torch_greater_or_equal("2.7"): + from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter + + +def _ensure_torch_distributed(device_type: str): + """Initialize torch.distributed if not already initialized.""" + if not torch.distributed.is_initialized(): + try: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + backend_map = { + "cuda": "nccl", + "cpu": "gloo", + "xpu": "xccl", + "hpu": "hccl", + "neuron": "neuron", + "tpu": "tpu_dist", + } + backend = backend_map.get(device_type) + + # Bind the accelerator before init so the process group is created with a + # device_id, otherwise collectives like barrier() warn (and may spin up an + # extra NCCL comm) about the missing device binding. + device_id = None + if device_type != "cpu": + getattr(torch, device_type).set_device(local_rank) + device_id = torch.device(device_type, local_rank) + torch.distributed.init_process_group( + backend=backend, rank=rank, world_size=world_size, device_id=device_id + ) + except Exception as e: + raise OSError( + "We tried to initialize torch.distributed for you, but it failed. Make " + "sure you init torch distributed in your script to use distributed training." + ) from e + + +def _distributed_barrier(): + """Barrier bound to the current accelerator device. + + Passing `device_ids` is required when the process group was initialized without a + `device_id`; with it, the call is a no-op compared to plain `barrier()`. Safe to call + when torch.distributed has not been initialized — returns immediately. + """ + if not torch.distributed.is_initialized(): + return + device_type = torch._C._get_accelerator().type + if device_type != "cpu": + torch.distributed.barrier(device_ids=[getattr(torch, device_type).current_device()]) + else: + torch.distributed.barrier() + + +def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed.device_mesh.DeviceMesh: + if not is_torch_greater_or_equal("2.5"): + raise OSError("Distributed training with DistributedConfig requires `torch>=2.5`.") + + if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.6"): + raise OSError("FSDP2 requires `torch>=2.6`.") + + device_type = torch._C._get_accelerator().type + _ensure_torch_distributed(device_type) + + world_size = torch.distributed.get_world_size() + if device_type != "cpu": + getattr(torch, device_type).set_device(int(os.environ.get("LOCAL_RANK", 0))) + + tp_size = distributed_config.tp_size + fsdp_size = distributed_config.fsdp_size + + assert world_size == tp_size * fsdp_size, ( + f"world_size ({world_size}) must be equal to tp_size ({tp_size}) * fsdp_size ({fsdp_size})" + ) + + dims, names = [], [] + if fsdp_size > 1: + dims.append(fsdp_size) + names.append("fsdp") + if tp_size > 1: + dims.append(tp_size) + names.append("tp") + + # Build the N-dimensional device mesh + mesh = torch.distributed.init_device_mesh(device_type, tuple(dims), mesh_dim_names=tuple(names)) + # If N > 1, create a flattened sub-mesh so all-reduces across the world mesh ae done in one collective + if len(dims) > 1: + mesh._flatten("_".join(names)) + + return mesh + + +def distribute_model(model, distributed_config: DistributedConfig, device_mesh) -> nn.Module: + """Apply TP and/or FSDP2 to `model` based on the mesh dims in `device_mesh`.""" + model.config.distributed_config = distributed_config + model.device_mesh = device_mesh + mesh_dim_names = device_mesh.mesh_dim_names or () + if "tp" in mesh_dim_names: + tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh + model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) + if "fsdp" in mesh_dim_names: + fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh + model = apply_fully_shard_data_parallel(model, fsdp_mesh, distributed_config.fsdp_plan) + return model + + +@torch.no_grad() +def clip_grad_norm(parameters, max_norm: float, norm_type: float = 2.0): + """Grad-norm clip that works when params live on different DTensor meshes. + + ``torch.nn.utils.get_total_norm`` stacks per-grad norms; that fails when grads + live on different meshes (e.g. TP-wrapped params on the (fsdp, tp) mesh and + FSDP-only params on the (fsdp,) sub-mesh). We sidestep it by replicating each + DTensor grad to a plain local tensor, computing the norm over those, and + scaling the original DTensor grads in place — the placement of the original + grads doesn't matter for the per-element clip. + """ + grads = [p.grad for p in parameters if p.grad is not None] + local_grads = [_replicate_dtensor(g).to_local() if isinstance(g, DTensor) else g for g in grads] + total_norm = torch.nn.utils.get_total_norm(local_grads, norm_type=norm_type) + torch.nn.utils.clip_grads_with_norm_(grads, max_norm=max_norm, total_norm=total_norm) + return total_norm + + +def gather_full_state_dict(model) -> dict[str, torch.Tensor]: + """Gather all sharded params to full plain tensors for saving. + + Handles FSDP unshard and TP DTensor gather. + Streams one parameter at a time to avoid holding all full tensors on GPU. + Only rank 0 accumulates the result; other ranks return ``{}``. + """ + is_rank0 = torch.distributed.get_rank() == 0 + state_dict = get_model_state_dict(model) + + result = {} + for key, tensor in state_dict.items(): + if not isinstance(tensor, DTensor): + if is_rank0: + result[key] = tensor.detach().to(device="cpu", copy=True).contiguous() + continue + + with torch.no_grad(): + full = _replicate_dtensor(tensor).to_local() + if is_rank0: + result[key] = full.detach().to(device="cpu", copy=True).contiguous() + del full + return result + + +def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None: + """Save model parameters as standard HF-format sharded safetensors using + DCP + HuggingFaceStorageWriter with consolidation enabled. + + Every rank first writes its own shard in parallel under + `/sharded/`, then a consolidation pass reads those shards + and emits HF-compatible `model-*-of-N.safetensors` (+ index) at + `/`. The result is a directory `from_pretrained` reads + through its normal path — no special flag needed at load time. + + DTensors carrying an uncomposed `_StridedShard` placement (e.g. fused + gate||up MoE weights) are replicated to a full tensor on every rank + before the save, otherwise DCP cannot encode that placement. + """ + if not is_torch_greater_or_equal("2.7"): + raise OSError("Distributed checkpoint saving requires `torch>=2.7`.") + + state_dict = get_model_state_dict(model) + for key, value in list(state_dict.items()): + if ( + isinstance(value, DTensor) + and _find_strided_shard_placement_from_fused_params(value.placements) is not None + ): + state_dict[key] = _replicate_dtensor(value) + + dcp.save( + state_dict, + storage_writer=HuggingFaceStorageWriter( + path=checkpoint_dir, + save_distributed=True, + enable_consolidation=True, + ), + ) + # Wait for rank 0 to finish writing the HF safetensors so other + # ranks don't return (and hit `from_pretrained`) before the files exist. + _distributed_barrier() + + +def has_mixed_tensor_and_dtensor_params(params) -> bool: + has_dtensor = False + has_tensor = False + for param in params: + if isinstance(param, DTensor): + has_dtensor = True + elif isinstance(param, torch.Tensor): + has_tensor = True + + if has_dtensor and has_tensor: + return True + return False + + +def maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) -> None: + """ + When get_optimizer_state_dict() or set_optimizer_state_dict() runs on an optimizer with no state yet, + PyTorch first materializes that state by doing a no-op step() with zero gradients. If an optimizer + group mixes regular tensors and DTensors, the batched foreach/fused optimizer kernels cannot process + that mixed group, so we turn those kernels off for such groups before distributed optimizer save/ + load. + """ + for i, param_group in enumerate(optimizer.param_groups): + if has_mixed_tensor_and_dtensor_params(param_group.get("params", ())): + logger.warning_once( + f"Param group {i} mixes regular tensors and DTensors; disabling foreach/fused " + "optimizer kernels for that group so distributed optimizer save/load can materialize state." + ) + param_group["foreach"] = False + if "fused" in param_group: + param_group["fused"] = False + + +def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: + """Save optimizer state via DCP. + + Params whose DTensors carry a lonely `_StridedShard` placement (e.g. Mixtral + `gate_up_proj`) are locally split into plain-`Shard` pieces at the boundary + so DCP only ever sees DTensors it can encode as one contiguous chunk per + rank. + """ + maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) + optimizer_state_dict = get_optimizer_state_dict(model, optimizer) + fusion_metadata = get_fusion_metadata(optimizer_state_dict) + unfuse_optimizer_state(optimizer_state_dict, fusion_metadata) + dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir) + + +def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None: + """Load optimizer state via DCP. + + Symmetric to `save_optimizer_distributed`: build the unfused template, let DCP fill + it, then merge fused params back to their original `_StridedShard` form + before handing the state_dict back to the optimizer. + """ + maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) + optimizer_state_dict = get_optimizer_state_dict(model, optimizer) + fusion_metadata = get_fusion_metadata(optimizer_state_dict) + unfuse_optimizer_state(optimizer_state_dict, fusion_metadata) + dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir) + fuse_optimizer_state(optimizer_state_dict, fusion_metadata) + set_optimizer_state_dict(model, optimizer, optimizer_state_dict) + maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 3a4875cdd990..aa5379f86b11 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -33,6 +33,7 @@ QuantizedCache, StaticCache, ) +from ..distributed.fsdp import is_fsdp_managed_module from ..dynamic_module_utils import ( check_python_requirements, get_cached_module_file, @@ -40,7 +41,6 @@ resolve_trust_remote_code, ) from ..integrations.deepspeed import is_deepspeed_zero3_enabled -from ..integrations.fsdp import is_fsdp_managed_module from ..masking_utils import create_masks_for_generate from ..tokenization_python import ExtensionsTrie from ..utils import ( @@ -2152,7 +2152,7 @@ def _extract_generation_mode_kwargs( "assistant_model": assistant_model, "streamer": streamer, } - world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 # type: ignore + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2599,7 +2599,7 @@ def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, # The following logic allows an early break if all peers finished generating their sequence this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) # send 0.0 if we finished, 1.0 otherwise - dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # type: ignore + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # did all peers finish? the reduced sum will be 0.0 then if this_peer_finished_flag.item() == 0.0: return False diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 336db3773f76..f6282cf4609e 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,6 @@ "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], - "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "ggml": [ "GGUF_CONFIG_DEFAULTS_MAPPING", "GGUF_CONFIG_MAPPING", @@ -160,11 +159,6 @@ "convert_and_export_with_cache", ] -_import_structure["tensor_parallel"] = [ - "shard_and_distribute_module", - "ALL_PARALLEL_STYLES", - "translate_to_torch_parallel_style", -] try: if not is_torch_greater_or_equal("2.5"): raise OptionalDependencyNotAvailable() @@ -209,7 +203,6 @@ from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear - from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .ggml import ( GGUF_CONFIG_DEFAULTS_MAPPING, GGUF_CONFIG_MAPPING, @@ -305,12 +298,6 @@ else: from .executorch import TorchExportableModuleWithStaticCache, convert_and_export_with_cache - from .tensor_parallel import ( - ALL_PARALLEL_STYLES, - shard_and_distribute_module, - translate_to_torch_parallel_style, - ) - try: if not is_torch_greater_or_equal("2.5"): raise OptionalDependencyNotAvailable() diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index 142489c4c21a..69df6a618bb2 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -26,6 +26,7 @@ from safetensors import safe_open from safetensors.torch import save_file +from ..distributed.fsdp import is_fsdp_enabled from ..utils import ( is_accelerate_available, is_torch_available, @@ -34,7 +35,6 @@ ) from ..utils.quantization_config import QuantizationMethod from .deepspeed import is_deepspeed_zero3_enabled -from .fsdp import is_fsdp_enabled if is_torch_available(): diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py deleted file mode 100644 index 7cda7ad55acc..000000000000 --- a/src/transformers/integrations/fsdp.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from __future__ import annotations - -import inspect -import os -from typing import TYPE_CHECKING - -from ..utils import is_torch_available, strtobool -from ..utils.quantization_config import QuantizationMethod - - -if TYPE_CHECKING: - from torch import nn - - -def is_fsdp_managed_module(module: nn.Module) -> bool: - if not is_torch_available(): - return False - - import torch - - if not torch.distributed.is_available(): - return False - - import torch.distributed.fsdp - - return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr( - module, "_is_fsdp_managed_module", False - ) - - -def is_fsdp_enabled(): - if is_torch_available(): - import torch - - return ( - torch.distributed.is_available() - and torch.distributed.is_initialized() - and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 - and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 - ) - - return False - - -def get_fsdp_ckpt_kwargs(): - """ - Returns checkpoint kwargs for FSDP model saving. - - Checks if the `adapter_only` parameter is supported by `save_fsdp_model` from accelerate - and returns the appropriate kwargs. - """ - from accelerate.utils import save_fsdp_model - - if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): - return {"adapter_only": True} - else: - return {} - - -def update_fsdp_plugin_peft(model, accelerator): - """ - Updates the FSDP plugin for PEFT LoRA/QLoRA compatibility. - - When using FSDP with PEFT LoRA, the auto wrap policy needs to be updated to additionally wrap - LoRA trainable layers separately. When using FSDP with QLoRA, the mixed precision policy needs - to be updated to use the quantization storage data type. - """ - from peft import PeftConfig - from peft.utils.other import fsdp_auto_wrap_policy - - if isinstance(model.active_peft_config, PeftConfig): - accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) - if ( - getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES - and model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point - ): - accelerator.state.fsdp_plugin.set_mixed_precision( - model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True - ) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 4f107f03e6fb..02369999c11f 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -366,7 +366,7 @@ def _grouped_linear( out = _grouped_mm(input, weight, offs=offs) else: # (S, input_dim) @ grouped (num_experts, output_dim, input_dim).T -> (S, output_dim) - out = _grouped_mm(input, weight.transpose(-2, -1), offs=offs) + out = _grouped_mm(input, weight.transpose(-2, -1).contiguous(), offs=offs) if bias is not None: # We should be able to pass bias to the grouped_mm call, but it's not yet supported. diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 67d9420659af..482ef891771f 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -511,28 +511,15 @@ def mlp_forward(self, hidden_states): def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs): - from ..integrations.tensor_parallel import shard_and_distribute_module - - model = kwargs.get("model") - empty_param = kwargs.get("empty_param") - casting_dtype = kwargs.get("casting_dtype") - to_contiguous = kwargs.get("to_contiguous") - rank = kwargs.get("rank") device_mesh = kwargs.get("device_mesh") + if device_mesh is not None: + raise NotImplementedError( + "MXFP4 quantization is not yet compatible with DTensor-based tensor parallelism. " + "Please disable TP or use a non-quantized model." + ) for proj in ["gate_up_proj", "down_proj"]: if proj in param_name: - if device_mesh is not None: - param_value = shard_and_distribute_module( - model, - param_value, - empty_param, - dq_param_name, - casting_dtype, - to_contiguous, - rank, - device_mesh, - ) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" setattr(module, param_name.rsplit(".", 1)[1], param_value) @@ -557,24 +544,19 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, trito triton_kernels_hub.matmul_ogs.FlexCtx, triton_kernels_hub.matmul_ogs.InFlexData, ) - from ..integrations.tensor_parallel import shard_and_distribute_module - model = kwargs.get("model") - empty_param = kwargs.get("empty_param") - casting_dtype = kwargs.get("casting_dtype") - to_contiguous = kwargs.get("to_contiguous") - rank = kwargs.get("rank") device_mesh = kwargs.get("device_mesh") + if device_mesh is not None: + raise NotImplementedError( + "MXFP4 quantization is not yet compatible with DTensor-based tensor parallelism. " + "Please disable TP or use a non-quantized model." + ) + if "blocks" in param_name: proj = param_name.split(".")[-1].split("_blocks")[0] if "scales" in param_name: proj = param_name.split(".")[-1].split("_scales")[0] - if device_mesh is not None: - shard_and_distribute_module( - model, param_value, empty_param, param_name, casting_dtype, to_contiguous, rank, device_mesh - ) - else: - setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) + setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) blocks_attr = f"{proj}_blocks" scales_attr = f"{proj}_scales" blocks = getattr(module, blocks_attr) # at this point values were loaded from ckpt diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py deleted file mode 100644 index 76335a2fa00d..000000000000 --- a/src/transformers/integrations/tensor_parallel.py +++ /dev/null @@ -1,1560 +0,0 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from __future__ import annotations - -import math -import operator -import os -import re -from functools import reduce - -from ..distributed import DistributedConfig -from ..utils import is_torch_greater_or_equal, logging -from ..utils.generic import GeneralInterface -from ..utils.import_utils import is_torch_available - - -if is_torch_available(): - import torch - import torch.distributed as dist - from torch import nn - - # Cache this result has it's a C FFI call which can be pretty time-consuming - _torch_distributed_available = torch.distributed.is_available() - - -logger = logging.get_logger(__name__) - - -def initialize_tensor_parallelism( - tp_plan: str | dict[str, str] | None, tp_size: int | None = None, device_mesh=None, device_map=None -): - r""" - Sets up the device mesh and initialized the backend for tensor parallelism. - This function is called when the model is loaded and the TP plan is set to 'auto'. - """ - if tp_size is not None and tp_plan is None: - raise ValueError("tp_plan has to be set when tp_size is passed.") - if tp_plan is not None and device_map is not None: - raise ValueError("`tp_plan` and `device_map` are mutually exclusive. Choose either one for parallelization.") - if device_mesh is None: - if not is_torch_greater_or_equal("2.5"): - raise OSError("Tensor parallel is only supported for `torch>=2.5`.") - - # Detect the accelerator on the machine. If no accelerator is available, it returns CPU. - device_type = torch._C._get_accelerator().type - if device_type == "mps": - raise RuntimeError("Tensor parallelism is not supported on MPS devices.") - current_device = getattr(torch, device_type) - if not torch.distributed.is_initialized(): - try: - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - - backend_map = { - "cuda": "nccl", - "cpu": "gloo", - "xpu": "xccl", - "hpu": "hccl", - "neuron": "neuron", - "tpu": "tpu_dist", - } - backend = backend_map.get(device_type) - - torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size) - current_device = getattr(torch, device_type) - if device_type != "cpu": - current_device.set_device(local_rank) - - except Exception as e: - raise OSError( - "We tried to initialize torch.distributed for you, but it failed. Make " - "sure you init torch distributed in your script to use `tp_plan`." - ) from e - - if device_type != "cpu": - current_device.set_device(int(os.environ["LOCAL_RANK"])) - index = current_device.current_device() - tp_device = torch.device(device_type, index) - device_map = tp_device - else: - tp_device = torch.device(device_type) - device_map = device_type or {} - - tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size() - device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,)) - else: - if device_mesh.ndim > 1: - if "tp" not in device_mesh.mesh_dim_names: - raise ValueError( - "When using `tp_plan` and n-d `device_mesh`, it must contain a 'tp' dimension. " - "Please provide a valid `device_mesh`." - ) - device_mesh = device_mesh["tp"] - tp_size = device_mesh.size() - device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}") - - return device_map, device_mesh, tp_size - - -def replace_layer_number_by_wildcard(name: str) -> str: - """ - Replace the numbers in the `name` by wildcards, only if they are in-between dots (`.`) or if they are between - a dot (`.`) and the end of the string. - This matches how modules are named/numbered when using a nn.ModuleList or nn.Sequential, but will NOT match - numbers in a parameter name itself, e.g. if the param is named `"w1"` or `"w2"`. - """ - return re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), name) - - -def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weight=True) -> str | None: - """ - Get the TP style for a parameter from the TP plan. - - The TP plan is a dictionary that maps parameter names to TP styles. - The parameter name can be a generic name with wildcards (e.g. "*.weight") or a specific name (e.g. "layer_1.weight"). - - The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but - not parent classes for `post_init` calls - """ - generic_param_name = replace_layer_number_by_wildcard(parameter_name) - if generic_param_name in tp_plan: - return tp_plan[generic_param_name] - elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: - return tp_plan[module_name] - return None - - -# ============================================================================= -# Tensor Sharding Utilities -# ============================================================================= - - -if is_torch_available(): - str_to_dtype = { - "BOOL": torch.bool, - "U8": torch.uint8, - "I8": torch.int8, - "I16": torch.int16, - "F16": torch.float16, - "BF16": torch.bfloat16, - "I32": torch.int32, - "F32": torch.float32, - "F64": torch.float64, - "I64": torch.int64, - "F8_E4M3": torch.float8_e4m3fn, - } - - -def _blocks_to_block_sizes(total_size: int, blocks: int | list[int]) -> list[int]: - """ - Convert block count or proportions to block sizes. - - This function accepts - - - The number of blocks (int), in which case the block size is - total_size//blocks; or - - A list of block sizes (list[int]). - - In the second case, if sum(blocks) < total_size, the ratios between - the block sizes will be preserved. For instance, if blocks is - [2, 1, 1] and total_size is 1024, the returned block sizes are - [512, 256, 256]. - """ - if isinstance(blocks, list): - total_blocks = sum(blocks) - assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}" - part_size = total_size // total_blocks - return [part_size * block for block in blocks] - else: - assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}" - single_size = total_size // blocks - return [single_size] * blocks - - -def get_packed_weights(param, empty_param, device_mesh, rank, dim): - """ - When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share. - So if you have: gate_proj ( 16, 5120, 8190) - and up_proj ( 16, 5120, 8190) - packed as gate_up_proj ( 16, 5120, 2 * 8190) - And you shard along the last dimension, you need to interleave the gate and up values: - - Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly. - - Let's take TP_size = 4 for an example: - - Packed tensor `gate_up_proj` - --------------------------------------------------------------- - [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ] - ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ - Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1 - - Explanation: - - The first half of the tensor (left of the center) holds the gate_proj values. - - The second half (right of the center) holds the up_proj values. - - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity. - - Each shard receives one slice from the gate part and the corresponding slice from the up part. - - For instance: - • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ] - • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ] - • … and so on. - - This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism. - """ - slice_ = param - total_size = empty_param.shape[dim] - world_size = device_mesh.size() - block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2) - - tensors_slices = [] - block_offset = 0 - for block_size in block_sizes: - shard_block_size = block_size // world_size - start = rank * shard_block_size - stop = (rank + 1) * shard_block_size - tensors_slices += range(block_offset + start, block_offset + stop) - block_offset += block_size - - slice_dtype = slice_.get_dtype() - # Handle F8_E4M3 dtype by converting to float16 before slicing - # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn' - casted = False - if slice_dtype == "F8_E4M3" or slice_dtype == "F8_E5M2": - slice_ = slice_[...].to(torch.float16) - casted = True - - if dim == 0: - tensor = slice_[tensors_slices, ...] - elif dim == 1 or dim == -2: - tensor = slice_[:, tensors_slices, ...] - elif dim == 2 or dim == -1: - tensor = slice_[..., tensors_slices] - else: - raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported") - - if casted: - return tensor - else: - return tensor.to(str_to_dtype[slice_dtype]) - - -def repack_weights( - packed_parameter: torch.Tensor, - sharded_dim: int, # The dimension index in the global tensor that was sharded - world_size: int, - num_blocks: int = 2, -) -> torch.Tensor: - """ - Reorders a tensor that was reconstructed from sharded packed weights into its canonical packed format. - - For example, if a weight was packed (e.g., gate_proj and up_proj) and then sharded, - DTensor.full_tensor() might produce an interleaved layout like [G0, U0, G1, U1, ...] - along the sharded dimension. This function reorders it to [G0, G1, ..., U0, U1, ...]. - This is an inverse operation to get_packed_weights. - - Args: - reconstructed_tensor: The tensor reconstructed from DTensor (e.g., via .full_tensor().contiguous()). - sharded_dim: The dimension index in the reconstructed_tensor that was originally sharded. - world_size: The tensor parallel world size. - num_packed_projs: The number of projections that were packed together (e.g., 2 for gate_up_proj). - - Returns: - The reordered tensor in canonical packed format. - """ - - if num_blocks != 2: - raise ValueError( - "Num blocks different from 2 is not supported yet. This is most likely a bug in your implementation as we only pack gate and up projections together." - ) - - actual_sharded_dim = sharded_dim if sharded_dim >= 0 else sharded_dim + packed_parameter.ndim - total_size_on_sharded_dim = packed_parameter.shape[actual_sharded_dim] - original_block_size_on_dim = total_size_on_sharded_dim // num_blocks - shard_chunk_size = original_block_size_on_dim // world_size - - prefix_shape = packed_parameter.shape[:actual_sharded_dim] - suffix_shape = packed_parameter.shape[actual_sharded_dim + 1 :] - - tensor_view = packed_parameter.view( - *prefix_shape, - world_size, - num_blocks, - shard_chunk_size, - *suffix_shape, - ) - - # Permute to bring num_packed_projs first, then world_size, then shard_chunk_size - # This groups all chunks of G together, then all chunks of U together. - # Target order of these middle dimensions: (num_packed_projs, world_size, shard_chunk_size) - # Current order of view's middle dimensions: (world_size, num_packed_projs, shard_chunk_size) - # Absolute indices of the dimensions to be permuted (world_size, num_packed_projs) - axis_ws_abs = len(prefix_shape) - axis_npp_abs = len(prefix_shape) + 1 - - permute_order = list(range(tensor_view.ndim)) - permute_order[axis_ws_abs], permute_order[axis_npp_abs] = permute_order[axis_npp_abs], permute_order[axis_ws_abs] - - tensor_permuted = tensor_view.permute(*permute_order) - - # Reshape back to the original tensor's ndim, with the sharded dimension now correctly ordered as [G_all, U_all]. - # The final shape should be the same as reconstructed_tensor. - final_ordered_tensor = tensor_permuted.reshape_as(packed_parameter) - - return final_ordered_tensor - - -def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: int | None = None): - """ - Generalized tensor sharding across a multi-dimensional device mesh. - Extract only the fraction of the parameter owned by the given `rank` when the parameter would have gone sharding at provided `dim`. - Extraction follows the pytorch `Shard` placement so that sharding and materializing back to full tensor follows `Shard` semantics. - `Shard` follows torch.chunk style sharding of the tensor. We demonstrate some cases below on how sharding happens including some edge cases - such as some ranks having an empty tensor as shard. Below implementation is robut to all these cases. - - Case (1) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 4 - rank 0 gets (4, 5120, 8190) (0 ... 4, 5120, 8190) - rank 1 gets (4, 5120, 8190) (4 ... 8, 5120, 8190) - rank 2 gets (4, 5120, 8190) (8 ... 12, 5120, 8190) - rank 3 gets (4, 5120, 8190) (12 ... 16, 5120, 8190) - - Case (2) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 14 - rank 0 gets (2, 5120, 8190) (0 ... 2, 5120, 8190) - rank 1 gets (2, 5120, 8190) (2 ... 4, 5120, 8190) - rank 2 gets (2, 5120, 8190) (4 ... 6, 5120, 8190) - rank 3 gets (2, 5120, 8190) (6 ... 8, 5120, 8190) - rank 4 gets (2, 5120, 8190) (8 ... 10, 5120, 8190) - rank 5 gets (2, 5120, 8190) (10 ... 12, 5120, 8190) - rank 6 gets (2, 5120, 8190) (12 ... 14, 5120, 8190) - rank 7 gets (2, 5120, 8190) (14 ... 16, 5120, 8190) - rank 8 gets (0, 5120, 8190) - rank 9 gets (0, 5120, 8190) - rank 10 gets (0, 5120, 8190) - rank 11 gets (0, 5120, 8190) - rank 12 gets (0, 5120, 8190) - rank 13 gets (0, 5120, 8190) - - Case (3) - empty_param (16, 5120, 8190) - dim 0 - device_mesh.size() 3 - rank 0 gets (6, 5120, 8190) (0 ... 6, 5120, 8190) - rank 1 gets (6, 5120, 8190) (6 ... 12, 5120, 8190) - rank 2 gets (4, 5120, 8190) (12 ... 16, 5120, 8190) - - In case (2), empty shards are returned with appropriate dimension to allow for operations to work smoothly. - Args: - param (torch.Tensor): The tensor to shard. - empty_param (torch.Tensor): A tensor used for shape reference. - device_mesh (torch.Tensor): Shape [d_0, ..., d_n] representing the mesh. - rank (int): Global rank of the current process/device. - dim (int): Dimension along which to shard the tensor. - """ - param_dim = empty_param.ndim - mesh_shape = device_mesh.shape - world_size = reduce(operator.mul, mesh_shape) - # Get param shape: works for both torch.Tensor and safetensors TensorInfo - param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() - if dim < 0: - dim = param_dim + dim - if empty_param.dim() == 3 and dim == 1 and len(param_shape) == 2: - dim = 0 - elif empty_param.dim() == 3 and dim == 2 and len(param_shape) == 2: - dim = 1 - - shard_size = math.ceil(param_shape[dim] / world_size) - start = rank * shard_size - end = min(start + shard_size, param_shape[dim]) - - if dim >= param_dim: - raise ValueError(f"dim {dim} is out of bounds for tensor of dimension {param_dim}") - - if rank >= world_size: - raise ValueError(f"Rank {rank} is out of bounds for mesh size {world_size}") - - # we have the full tensor not 1 part of it. - # in that case, we just assume that the weight was properly saved - # and thus because we TP if the layer is colwise it should not use this. Layer should be packed_colwise - # to inform that it needs to read form a packed tensor. It will also take care of the module list thingy. - # here we take care of potential chunking / layer split / layer chunking. - # The only "hard" case is? if we collect q,k,v -> merge it into qkv. In that case - # actually we still shard dim=0 does not change - # so only case is if the dim of the empty param is 3 and the shard dim is 0 -> we put the - # tensor on a certain device (with the input tensor_index) - if tensor_idx is not None and empty_param.dim() == 3 and dim == 0 and len(param_shape) == 2: - # special case we don't "shard" just send this entire tensor to the correct rank. - if start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:] - else: - return torch.empty([], dtype=torch.int64, device=rank) - - slice_indices = [slice(None)] * len(param_shape) - - if start < param_shape[dim]: - slice_indices[dim] = slice(start, end) - param = param[tuple(slice_indices)] - if isinstance(param, list): # TODO handle the modulelist case! - param = [p[:] for p in param] - return param - - param_shape[dim] = 0 - return torch.empty(tuple(param_shape), dtype=torch.int64) # empty allocates memory.... - - -def _split_along_last_dim(x, world_size): - """Split tensor along last dimension into world_size chunks.""" - return torch.chunk(x, world_size, dim=-1) - - -# ============================================================================= -# Distributed Communication Primitives -# ============================================================================= -# -# Naming convention: -# - Functions describe their FORWARD behavior -# - Backward behavior is the "conjugate" operation for gradient flow -# -# Available operations: -# ┌────────────────────┬─────────────────────┬─────────────────────┐ -# │ Function │ Forward │ Backward │ -# ├────────────────────┼─────────────────────┼─────────────────────┤ -# │ all_reduce │ all-reduce (sum) │ identity │ -# │ all_reduce_backward│ identity │ all-reduce (sum) │ -# │ all_gather │ all-gather │ split (local chunk) │ -# │ split │ split (local chunk) │ all-gather │ -# │ reduce_scatter │ reduce-scatter │ all-gather │ -# └────────────────────┴─────────────────────┴─────────────────────┘ -# =================== - - -class _AllReduceBackward(torch.autograd.Function): - """Identity forward, all-reduce backward. Used before colwise layers (f in Megatron).""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - return x - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - if device_mesh.size() == 1: - return grad_output, None - grad_output = grad_output.contiguous() - dist.all_reduce(grad_output, op=dist.ReduceOp.SUM, group=device_mesh.get_group()) - return grad_output, None - - -class _AllReduceForward(torch.autograd.Function): - """All-reduce forward, identity backward. Used after rowwise layers (g in Megatron).""" - - @staticmethod - def forward(ctx, x, device_mesh): - if device_mesh.size() == 1: - return x - dist.all_reduce(x, op=dist.ReduceOp.SUM, group=device_mesh.get_group()) - return x - - @staticmethod - def backward(ctx, grad_output): - return grad_output, None - - -class _AllGather(torch.autograd.Function): - """All-gather forward, split backward. Gathers sharded outputs.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - last_dim = x.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - x = x.contiguous() - tensor_list = [torch.empty_like(x) for _ in range(world_size)] - tensor_list[rank] = x - dist.all_gather(tensor_list, x, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous() - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - rank = device_mesh.get_local_rank() - chunks = _split_along_last_dim(grad_output, world_size) - return chunks[rank].contiguous(), None - - -class _Split(torch.autograd.Function): - """Split forward, all-gather backward. Scatters replicated input.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - rank = device_mesh.get_local_rank() - chunks = _split_along_last_dim(x, world_size) - return chunks[rank].contiguous() - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - last_dim = grad_output.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - grad_output = grad_output.contiguous() - tensor_list = [torch.empty_like(grad_output) for _ in range(world_size)] - tensor_list[rank] = grad_output - dist.all_gather(tensor_list, grad_output, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous(), None - - -class _ReduceScatter(torch.autograd.Function): - """Reduce-scatter forward, all-gather backward. For sequence parallel.""" - - @staticmethod - def forward(ctx, x, device_mesh): - ctx.device_mesh = device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return x - - last_dim = x.dim() - 1 - group = device_mesh.get_group() - - input_chunks = list(x.chunk(world_size, dim=last_dim)) - output_shape = list(x.shape) - output_shape[last_dim] //= world_size - output = torch.empty(output_shape, dtype=x.dtype, device=x.device) - - dist.reduce_scatter(output, input_chunks, op=dist.ReduceOp.SUM, group=group) - return output - - @staticmethod - def backward(ctx, grad_output): - device_mesh = ctx.device_mesh - world_size = device_mesh.size() - - if world_size == 1: - return grad_output, None - - last_dim = grad_output.dim() - 1 - rank = device_mesh.get_local_rank() - group = device_mesh.get_group() - - grad_output = grad_output.contiguous() - tensor_list = [torch.empty_like(grad_output) for _ in range(world_size)] - tensor_list[rank] = grad_output - dist.all_gather(tensor_list, grad_output, group=group) - return torch.cat(tensor_list, dim=last_dim).contiguous(), None - - -# ============================================================================= -# Convenience wrappers -# ============================================================================= - - -def all_reduce_backward(x, device_mesh): - """Identity forward, all-reduce backward. Use before colwise layers.""" - return _AllReduceBackward.apply(x, device_mesh) - - -def all_reduce_forward(x, device_mesh): - """All-reduce forward, identity backward. Use after rowwise layers.""" - return _AllReduceForward.apply(x, device_mesh) - - -def all_gather(x, device_mesh): - """All-gather forward, split backward.""" - return _AllGather.apply(x, device_mesh) - - -def split(x, device_mesh): - """Split forward, all-gather backward.""" - return _Split.apply(x, device_mesh) - - -def reduce_scatter(x, device_mesh): - """Reduce-scatter forward, all-gather backward.""" - return _ReduceScatter.apply(x, device_mesh) - - -def distribute_module( - module: nn.Module, - device_mesh=None, - input_fn=None, - output_fn=None, -) -> nn.Module: - """ - Copy pasted from torch's function but we remove the communications (partitioning) - as well as buffer registering that is similarly not efficient. - """ - if input_fn is not None: - module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh)) - if output_fn is not None: - module.register_forward_hook(lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)) - return module - - -class TensorParallelLayer: - """General tensor parallel layer for transformers""" - - device_mesh = None - rank = None - empty_param = None - - def __init__(self, device_mesh=None, rank=None, empty_param=None): - self.rank = rank - self.device_mesh = device_mesh - self.empty_param = empty_param - - def _prepare_input_fn(self, mod, inputs, device_mesh): - raise NotImplementedError - - def _prepare_output_fn(self, mod, outputs, device_mesh): - raise NotImplementedError - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - raise NotImplementedError - - def prepare_module_tp(self, module: nn.Module, device_mesh, **kwargs) -> nn.Module: - distribute_module( - module, - device_mesh, - self._prepare_input_fn, - self._prepare_output_fn, - ) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - """ - Compute the expected shape after TP sharding for a given full shape. - - Args: - full_shape: The full (unsharded) parameter shape - - Returns: - The expected sharded shape for this rank - """ - # Default: no sharding, return full shape - return tuple(full_shape) - - def update_module_attributes(self, module: nn.Module): - """ - Update module attributes (e.g. in_features, out_features) to reflect sharded dimensions. - - Args: - module: The module to update - - Returns: - None, update the module in-place - """ - pass - - -class ColwiseParallel(TensorParallelLayer): - """ - Column-wise parallel: weight is sharded on dim -2 (output features). - Forward: input replicated -> output sharded on last dim. - If gather_output=True, output is all-gathered to produce full tensor. - """ - - def __init__(self, gather_output: bool = False, **kwargs): - super().__init__(**kwargs) - self.gather_output = gather_output - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - return all_reduce_backward(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - if self.gather_output: - return all_gather(outputs, device_mesh) - return outputs - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # Colwise shards dim -2, but 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else -2 - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - # If we gather the output, the output dimension of the module is not sharded, so no need to update out_features. - # Otherwise, we need to update out_features to reflect the sharded dimension. - if not self.gather_output and hasattr(module, "out_features"): - module.out_features = self.get_expected_sharded_shape((module.out_features,))[0] - - -class ReplicatedWithGradAllReduce(TensorParallelLayer): - """ - Replicated parameter with gradient all-reduce. - - For parameters like q_norm/k_norm that sit between colwise and rowwise - layers. The parameter is replicated (not sharded), but its gradient - accumulates from local heads only in TP mode. This class registers a - backward hook to all-reduce the parameter gradient. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return outputs - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, **kwargs): - # Use a module-level backward hook (not param.register_hook) because parameters are replaced during weight loading after this method runs. - # Module hooks survive parameter replacement. - def _backward_hook(mod, grad_input, grad_output, mesh=device_mesh): - for param in mod.parameters(): - if param.grad is not None: - all_reduce_forward(param.grad, mesh) - - module.register_full_backward_hook(_backward_hook) - - -class MlaKvAProjParallel(TensorParallelLayer): - """ - For MLA attention used in DeepSeek-V2 style models (deepseek_v2, longcat_flash, glm_moe_dsa, glm4_moe_lite): - kv_a_proj_with_mqa output is [kv_lora_rank + qk_rope_head_dim] (can have different naming but important thing - to understand is that it is split) - Example below (from modeling_longcat_flash.py): - - kv_a_proj_with_mqa - | - split - / \ - k_pass k_rot <-- "bypasses kv_b_proj" - | | (goes straight to attention, - kv_a_layernorm | never touches kv_b_proj) - | | - kv_b_proj | - (colwise) | - | | - k_pass k_rot - \\ / - cat - | - key_states - - k_pass is passed to kv_b_proj (colwise) which has built-in all_reduce_backward so we don't have a partial gradient for it. - However, k_rot goes straight to attention, never touches kv_b_proj. So we need to average gradient across all ranks otherwise we only get gradient for one rank (partial gradient). - """ - - def _prepare_output_fn(self, mod, output, device_mesh): - if not hasattr(mod.config, "qk_rope_head_dim"): - raise AttributeError( - f"Config for {type(mod).__name__} does not have `qk_rope_head_dim`. " - "MlaKvAProjParallel requires `qk_rope_head_dim` to be defined in the model config. " - "Please add it to the model's config or update the TP plan mapping." - ) - rope_dim = mod.config.qk_rope_head_dim - pass_output, rope_output = output.split([output.shape[-1] - rope_dim, rope_dim], dim=-1) - rope_output = all_reduce_backward(rope_output, device_mesh) - return torch.cat([pass_output, rope_output], dim=-1) - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, config=None, **kwargs): - module.config = config - distribute_module(module, device_mesh, output_fn=self._prepare_output_fn) - - -class RowwiseParallel(TensorParallelLayer): - """ - Row-wise parallel: weight is sharded on dim -1 (input features). - Forward: input (optionally split) -> output partial -> all-reduce to replicate. - - Args: - split_input: If True, splits replicated input before matmul. Use when input - comes from a non-parallelizable operation (chunk/slice). - Default False (expects pre-sharded input from colwise layer). - """ - - def __init__(self, split_input: bool = False, **kwargs): - super().__init__(**kwargs) - self.split_input = split_input - - def _prepare_input_fn(self, mod, inputs, device_mesh): - if hasattr(mod, "bias") and mod.bias is not None: - mod._bias = mod.bias - mod.bias = None - - input_tensor = inputs[0] if inputs else inputs - - if self.split_input: - # Input is replicated, split it to match sharded weight - return split(input_tensor, device_mesh) - return input_tensor - - def _prepare_output_fn(self, mod, outputs, device_mesh): - outputs = all_reduce_forward(outputs, device_mesh) - if hasattr(mod, "_bias") and mod._bias is not None: - outputs = outputs + mod._bias - return outputs - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, it should not be sharded (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = param[...] - else: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # 1D tensors (bias) are NOT sharded in rowwise - if len(full_shape) == 1: - return tuple(full_shape) - world_size = self.device_mesh.size() - shape = list(full_shape) - dim = -1 - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "in_features"): - # To fall in the 2D case in get_expected_sharded_shape, - # otherwise it will be treated as 1D and not sharded - shape = (1, module.in_features) - module.in_features = self.get_expected_sharded_shape(shape)[1] - - -class PackedColwiseParallel(ColwiseParallel): - """Packed column-wise parallel for fused weights like gate_up_proj.""" - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - expected_shape = self.get_expected_sharded_shape(self.empty_param.shape) - if dim < len(expected_shape): - # Input is unpacked (e.g., gate_proj that will be concatenated to gate_up_proj) - # Use regular tensor shard - concatenation will happen after - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2) - else: - # Input is already packed, use packed sharding - parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -2) - return parameter.to(device=device, dtype=dtype) - - -class PackedRowwiseParallel(RowwiseParallel): - """Packed row-wise parallel for fused weights like gate_up_proj.""" - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, it should not be sharded (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = param[...] - else: - # Check if input tensor is unpacked (shape mismatch with expected packed size) - # This happens when using MergeModulelist + Concatenate for fused weights like gate_up_proj - param_shape = param.shape if isinstance(param, torch.Tensor) else param.get_shape() - expected_packed_dim = self.empty_param.shape[-1] if self.empty_param.dim() >= 1 else 0 - actual_dim = param_shape[-1] if len(param_shape) >= 1 else 0 - - if actual_dim < expected_packed_dim: - # Input is unpacked, use regular tensor shard - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - # Input is already packed, use packed sharding - parameter = get_packed_weights(param, self.empty_param, self.device_mesh, self.rank, -1) - return parameter.to(device=device, dtype=dtype) - - -class EmbeddingParallel(TensorParallelLayer): - """EmbeddingParallel: shards embedding table, handles masked lookups for vocab parallelism.""" - - def __init__(self, *, embedding_dim_sharding: int = 0, **kwargs): - super().__init__(**kwargs) - self.embedding_dim_sharding = embedding_dim_sharding - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - - # For vocab-parallel (dim 0), we need to handle masking and offsetting - if self.embedding_dim_sharding == 0: - rank = device_mesh.get_local_rank() - - # Get vocab range for this rank - # Use weight.shape[0] to get the actual local (sharded) size, not num_embeddings - # which may not be updated after sharding - per_partition_size = mod.weight.shape[0] - vocab_start_index = rank * per_partition_size - vocab_end_index = vocab_start_index + per_partition_size - - # Build mask for out-of-vocabulary tokens - input_mask = (input_tensor < vocab_start_index) | (input_tensor >= vocab_end_index) - mod._input_mask = input_mask - - # Offset input to local indices and mask invalid ones - masked_input = input_tensor.clone() - vocab_start_index - masked_input[input_mask] = 0 # Set to valid local index - - return masked_input - - return input_tensor - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # For vocab-parallel (dim 0), zero out embeddings for out-of-range tokens before all-reduce - if self.embedding_dim_sharding == 0 and hasattr(mod, "_input_mask"): - input_mask = mod._input_mask - # Use multiplication instead of in-place assignment to preserve gradients - mask_expanded = input_mask.unsqueeze(-1).expand_as(outputs) - outputs = outputs * (~mask_expanded).to(outputs.dtype) - del mod._input_mask - - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # If only 1 dim, shard this one (usually it's a `bias`) - dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape()) - if dim == 1: - parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1) - else: - parameter = get_tensor_shard( - param, - self.empty_param, - self.device_mesh, - self.rank, - self.embedding_dim_sharding, - ) - return parameter.to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - world_size = self.device_mesh.size() - shape = list(full_shape) - # EmbeddingParallel shards on self.embedding_dim_sharding (default 0) - # 1D tensors (bias) shard on dim -1 - dim = -1 if len(shape) == 1 else self.embedding_dim_sharding - dim = len(shape) + dim if dim < 0 else dim - shard_size = math.ceil(shape[dim] / world_size) - start = self.rank * shard_size - end = min(start + shard_size, shape[dim]) - shape[dim] = end - start - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_embeddings") and self.embedding_dim_sharding == 0: - module.num_embeddings = self.get_expected_sharded_shape((module.num_embeddings,))[0] - if hasattr(module, "embedding_dim") and self.embedding_dim_sharding == 1: - module.embedding_dim = self.get_expected_sharded_shape((module.embedding_dim,))[0] - - -class SequenceParallel(TensorParallelLayer): - """ - Sequence Parallel: input/output sharded on sequence dimension. - Weights are replicated. - """ - - def __init__(self, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False, **kwargs): - super().__init__(**kwargs) - self.sequence_dim = sequence_dim - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # For sequence parallel, input is sharded on sequence dim - # All-gather for the layer, then reduce-scatter after - return all_gather(input_tensor, device_mesh) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return reduce_scatter(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class GroupedGemmParallel(TensorParallelLayer): - """ - Applies Expert Parallelism to MoE experts by loading the correct experts on each device. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - global_num_experts = self.empty_param.shape[0] - if global_num_experts % self.device_mesh.size() != 0: - raise ValueError( - f"Global number of experts must be divisible by number of devices: {global_num_experts} % {self.device_mesh.size()} != 0" - ) - local_num_experts = global_num_experts // self.device_mesh.size() - shard_size = local_num_experts - start = self.rank * shard_size - end = (self.rank + 1) * shard_size - # special case we don't "shard" just send this entire tensor to the correct rank. - shape = param.get_shape() if not isinstance(param, torch.Tensor) else param.shape - if tensor_idx is not None and start <= tensor_idx < end: - # this tensor does need to be materialized on this device: - return param[:].to(device=device) - elif tensor_idx is None: # a bias or a weight, but already merged - return param[start:end].to(device=device, dtype=dtype) - elif len(shape) >= 1 and tensor_idx is not None: - return None - else: # bias case - return param[:].to(device=device, dtype=dtype) - - def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: - # GroupedGemm shards on dim 0 (experts dimension) - world_size = self.device_mesh.size() - shape = list(full_shape) - local_num_experts = shape[0] // world_size - shape[0] = local_num_experts - return tuple(shape) - - def update_module_attributes(self, module: nn.Module): - if hasattr(module, "num_experts"): - module.num_experts = self.get_expected_sharded_shape((self.empty_param.shape[0],))[0] - - -class RouterParallel(TensorParallelLayer): - """ - Allows to reshape the router scores to support running expert parallel. - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - return inputs[0] if inputs else inputs - - def _prepare_output_fn(self, mod, outputs, device_mesh): - """ - Remap global expert indices to local and zero out non-local scores. - - Example: 4 tokens, top_k=4, 128 experts, EP=8. num_local_experts = 128/8 = 16. - - Router produces (all ranks see the same values): - router_scores: (4, 4) — top-k routing weights - router_indices: (4, 4) — global expert IDs - [ 52, 42, 119, 67], - [102, 89, 61, 40], - [ 82, 103, 4, 34], - [ 93, 23, 109, 11], - - Each index maps to a rank: index // 16 gives the owning rank. - [3, 2, 7, 4], - [6, 5, 3, 2], - [5, 6, 0, 2], - [5, 1, 6, 0], - - For rank 0 (owns experts 0-15), we remap local indices with fmod and - fill non-local with sentinel=16 (used for one_hot masking): - router_indices (rank 0): - [ 16, 16, 16, 16], - [ 16, 16, 16, 16], - [ 16, 16, 4, 16], - [ 16, 16, 16, 11], - - Scores for non-local experts are zeroed out via masked_fill: - router_scores (rank 0): - [0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 0.3, 0.0], ← only expert 4 (local) keeps its score - [0.0, 0.0, 0.0, 0.1], ← only expert 11 (local) keeps its score - - both router_scores and router_indices stay (seq, top_k) shape. - They are paired element-wise: scores[i] is the weight for indices[i]. - All expert forward implementations (grouped_mm, batched_mm, eager) flatten - both with reshape(-1) and rely on this pairing. Changing the shape of one - without the other breaks routing! - - Each rank believes it is alone and computes only its part of the hidden states. - The sentinel index (num_local_experts) is skipped by one_hot encoding or clamped - + masked in grouped_mm/batched_mm. After the expert forward, an all_reduce sums - partial outputs across EP ranks to produce the full result. - """ - ep_rank, ep_size = device_mesh.get_local_rank(), device_mesh.size() - num_experts = getattr(mod, "num_experts", None) - if num_experts is None: - num_experts = getattr(getattr(mod, "config", None), "num_experts", None) - if num_experts is None: - raise AttributeError(f"Router module {type(mod).__name__} is missing num_experts and config.num_experts") - - if num_experts % ep_size != 0: - raise ValueError( - f"The number of experts must be divisible by number of ep_size: {num_experts} % {ep_size} != 0" - ) - num_local_experts = num_experts // ep_size - router_logits, router_scores, router_indices = outputs - non_local_mask = (router_indices // num_local_experts) != ep_rank - router_scores = router_scores.masked_fill(non_local_mask, 0.0) - router_indices = router_indices.masked_fill(non_local_mask, -1) - # As -1 % 1 is 0, we can only use mask fill when num_local_experts is 1 - if num_local_experts > 1: - router_indices = torch.fmod(router_indices, num_local_experts) - else: - router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1) - router_indices = router_indices.masked_fill(router_indices == -1, num_local_experts) - return router_logits, router_scores, router_indices - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - return param[...].to(device=device, dtype=dtype) - - -class MoeTensorParalellExperts(TensorParallelLayer): - """ - Note: For tensor parallel, the MoEExpertsParallel TP layer handles gradient sync: - - all_reduce_backward on hidden_states (for colwise gate_up_proj gradient) - - all_reduce_backward on top_k_weights (for router gradient) - - all_reduce_forward on output (for partial expert outputs) - """ - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def _prepare_input_fn(self, mod, inputs, device_mesh): - # inputs = (hidden_states, top_k_index, top_k_weights) - hidden_states = inputs[0] - top_k_index = inputs[1] - top_k_weights = inputs[2] - - # all_reduce_backward on hidden_states for correct colwise (gate_up_proj) gradient - hidden_states = all_reduce_backward(hidden_states, device_mesh) - - # all_reduce_backward on routing weights for correct router gradient - # This is needed because ∂L/∂routing_weights = ∂L/∂output * partial_expert_output - # and partial_expert_output is different on each GPU before all-reduce - top_k_weights = all_reduce_backward(top_k_weights, device_mesh) - - return (hidden_states, top_k_index, top_k_weights) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - # all_reduce_forward to sum partial expert outputs across GPUs - return all_reduce_forward(outputs, device_mesh) - - def shard_tensor( - self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None - ) -> torch.Tensor: - # This class doesn't shard tensors - sharding is handled by packed_colwise/rowwise - # on the individual weight tensors (gate_up_proj/down_proj) - return param[...].to(device=device, dtype=dtype) - - -class MoeIdentityExpertParallel(TensorParallelLayer): - """ - TP class for zero/identity experts in MoE layers. - - Under TP, the parent MoeTensorParalellExperts does all_reduce_forward (sum) - on the expert module output. Identity experts produce the same output on - every rank, so the sum gives world_size * output. This class divides the - input by world_size to compensate. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - input_tensor = inputs[0] if inputs else inputs - # TODO(fmom): when 2D-device mesh, need to select a //-ism axis to divide the input tensor by. - return input_tensor / device_mesh.size() - - def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None): - return param[...].to(device=device, dtype=dtype) - - def prepare_module_tp(self, module, device_mesh, **kwargs): - distribute_module(module, device_mesh, input_fn=self._prepare_input_fn) - - -class ParallelInterface(GeneralInterface): - # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if - # a new instance is created (in order to locally override a given entry) - _global_mapping = ( - { - "embedding_rowwise": EmbeddingParallel(embedding_dim_sharding=0), - "embedding_colwise": EmbeddingParallel(embedding_dim_sharding=1), - "colwise_gather_output": ColwiseParallel(gather_output=True), - "colwise": ColwiseParallel(), - "rowwise": RowwiseParallel(), - "rowwise_split_input": RowwiseParallel(split_input=True), - "packed_colwise": PackedColwiseParallel(), - "packed_rowwise": PackedRowwiseParallel(), - "sequence_parallel": SequenceParallel(), - "grouped_gemm": GroupedGemmParallel(), - "ep_router": RouterParallel(), - "moe_tp_experts": MoeTensorParalellExperts(), - "moe_identity_expert": MoeIdentityExpertParallel(), - "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), - "mla_kv_a_proj": MlaKvAProjParallel(), - } - if is_torch_available() and _torch_distributed_available - else {} - ) - - # Map plan names to sharding dimensions for weights - # For weights: colwise shards dim -2, rowwise shards dim -1 - # For embedding: rowwise shards dim 0 (vocab), colwise shards dim -2 (hidden) - plan_to_weight_dim: dict[str, int | None] = { - "colwise": -2, - "colwise_gather_output": -2, - "packed_colwise": -2, - "rowwise": -1, - "rowwise_split_input": -1, - "packed_rowwise": -1, - "embedding_rowwise": 0, - "embedding_colwise": 1, - "sequence_parallel": None, - "replicated_with_grad_allreduce": None, - "mla_kv_a_proj": None, - } - - # Bias sharding: colwise shards bias, rowwise doesn't (bias is replicated and all-reduced) - plan_to_bias_dim: dict[str, int | None] = { - "colwise": -1, - "colwise_gather_output": -1, - "packed_colwise": -1, - "rowwise": None, - "rowwise_split_input": None, - "packed_rowwise": None, - "embedding_rowwise": None, - "embedding_colwise": None, - "sequence_parallel": None, - "replicated_with_grad_allreduce": None, - "mla_kv_a_proj": None, - } - - @classmethod - def register_plan_to_weight_dim(cls, key: str, value: int | None): - cls.plan_to_weight_dim[key] = value - - @classmethod - def register_plan_to_bias_dim(cls, key: str, value: int | None): - cls.plan_to_bias_dim[key] = value - - -ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() - - -# ============================================================================= -# High-Level API Functions -# ============================================================================= - - -def gather_full_tensor( - local_tensor: torch.Tensor, shard_dim: int, device_mesh: dist.device_mesh.DeviceMesh -) -> torch.Tensor: - """ - All-gather a sharded tensor along the specified dimension to reconstruct the full tensor. - - Args: - local_tensor: The local shard of the tensor on this rank - shard_dim: The dimension along which the tensor was sharded - device_mesh: The device mesh for distributed communication - - Returns: - The full reconstructed tensor (same on all ranks) - """ - world_size = device_mesh.size() - # In case of TP+DP configuration, the TP group should be used for gathering, not the full DP group - process_group = device_mesh.get_group("tp") if "tp" in (device_mesh.mesh_dim_names or {}) else None - - # Normalize negative dimension - if shard_dim < 0: - shard_dim = local_tensor.ndim + shard_dim - - # Gather all shards - gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)] - dist.all_gather(gathered_tensors, local_tensor.contiguous(), group=process_group) - - # Concatenate along the shard dimension - return torch.cat(gathered_tensors, dim=shard_dim) - - -def gather_state_dict_for_save( - state_dict: dict[str, torch.Tensor], - tp_plan: dict[str, str], - device_mesh, - tp_size: int, -) -> dict[str, torch.Tensor]: - """ - Gather sharded tensors to reconstruct full tensors for saving. - - This function all-gathers each sharded tensor along its shard dimension - to reconstruct the full unsharded tensor for checkpoint saving. - - Args: - state_dict: The model state dict with local sharded tensors - tp_plan: The tensor parallel plan mapping layer patterns to shard styles - device_mesh: The device mesh for distributed communication - tp_size: The tensor parallel world size - - Returns: - State dict with full (gathered) tensors - """ - # Use the global mappings from ParallelInterface (can be extended by users) - plan_to_weight_dim = ALL_PARALLEL_STYLES.plan_to_weight_dim - plan_to_bias_dim = ALL_PARALLEL_STYLES.plan_to_bias_dim - - result = {} - for key, tensor in state_dict.items(): - # Find the matching TP plan for this parameter - param_name = key.rsplit(".", 1)[0] if "." in key else key - param_type = key.rsplit(".", 1)[1] if "." in key else None - generic_param_name = re.sub(r"\d+", "*", param_name) - # Also check the full key for nn.Parameter (e.g., MoE experts without .weight suffix) - generic_full_key = re.sub(r"\d+", "*", key) - - # Check if this parameter has a TP plan - current_plan = None - if generic_full_key in tp_plan: - # Full key match (e.g., "model.layers.*.mlp.experts.gate_up_proj" for MoE experts) - current_plan = tp_plan[generic_full_key] - elif generic_param_name in tp_plan: - current_plan = tp_plan[generic_param_name] - elif "." in generic_param_name: - parent_param_name = generic_param_name.rsplit(".", 1)[0] - if parent_param_name in tp_plan: - current_plan = tp_plan[parent_param_name] - - if current_plan is None or current_plan not in plan_to_weight_dim: - # Not sharded, keep as-is - result[key] = tensor - continue - - # Determine sharding dimension based on param type - if param_type == "bias": - shard_dim = plan_to_bias_dim.get(current_plan) - else: - shard_dim = plan_to_weight_dim.get(current_plan) - - if shard_dim is None: - # Replicated, keep as-is - result[key] = tensor - continue - - # Gather full tensor and handle packed weights repacking - full_tensor = gather_full_tensor(tensor, shard_dim, device_mesh) - if current_plan in ("packed_colwise", "packed_rowwise"): - full_tensor = repack_weights(full_tensor, shard_dim, tp_size, 2) - result[key] = full_tensor.contiguous() - - return result - - -def add_tensor_parallel_hooks_to_module( - model, - module, - current_module_plan, - layer_name, - device_mesh, -): - r""" - This function is called in `PretrainedModel.post_init()`. It is responsible of adding hooks - to the modules of the `model`, based on the `PretrainedModel._tp_plan`. - - This is the place where we add the `pre_forward` and `post_forwards` hooks. These are defined - for each `TensorParallelLayer` as `_prepare_input_fn` and `_prepare_output_fn`. - - Args: - model (`PretrainedModel`): The model containing the modules. - module (`nn.Module`): The current module to which we want to add the hooks. - current_module_plan (`str` or `None`): The tensor parallel plan for the current module, if any. - layer_name (`str`): The qualified name of the current module. - device_mesh (`dist.device_mesh.DeviceMesh`): The device mesh for distributed communication. - - """ - if current_module_plan is not None: - tp_layer = ALL_PARALLEL_STYLES[current_module_plan] - try: - tp_layer.prepare_module_tp(module, device_mesh, config=model.config) - except NotImplementedError as e: - logger.warning( - f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP " - f"plan: {e}" - ) - - module._hf_tp_plan = current_module_plan - module._hf_device_mesh = device_mesh - module.__repr__ = lambda: f"{module.__repr__()}\nTP Plan: {current_module_plan}" - - -def shard_and_distribute_module( - model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh -): - r""" - This function is called in `from_pretrained` when loading a model's checkpoints. - It receives the pointer to the parameter (or the parameter itself) and takes care of "sharding". - All process run this function, so they just load the partition of the tensor that they require. - - Main uses cases: - - column / rowise parallelism, you just shard all the weights of the layer (weight and bias) - - packed layers: you slice the weights, then shard like above - - custom operation: - - you want to add an all-gather at the end of a local layer. - - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance) - - """ - param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name - tp_plan = model.tp_plan or {} - module_to_tp = model.get_submodule(param_name) - rank = int(rank) - current_shard_plan = _get_parameter_tp_plan(parameter_name, tp_plan) - - if dist.get_rank() == 0: - if current_shard_plan is None: - logger.info(f"Tensor sharding plan for {param_name} not found, using default 'replicate' plan.") - else: - logger.info(f"Tensor sharding plan for {param_name}: {current_shard_plan}") - - tp_layer = None - if current_shard_plan is not None: - try: - tp_layer = ALL_PARALLEL_STYLES[current_shard_plan] - tp_layer.empty_param = empty_param - tp_layer.device_mesh = device_mesh - tp_layer.rank = rank - param = tp_layer.shard_tensor(param, tensor_idx=None, dtype=param_casting_dtype, device=rank) - if is_contiguous: - param = param.contiguous() - except NotImplementedError as e: - print( - f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}" - ) - else: - param = param[:].to(param_casting_dtype) - - # SUPER IMPORTANT we have to use setattr - # otherwise loading is crazy slow - if not isinstance(param, torch.nn.Parameter): - param = torch.nn.Parameter(param, requires_grad=empty_param.is_floating_point()) - setattr(module_to_tp, param_type, param) - if tp_layer is not None: - tp_layer.update_module_attributes(module_to_tp) - return param - - -def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): - """ - Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. - """ - - if tp_plan is None: - return - - generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} - unsharded_layers = set(generic_keys) - unused_rules = tp_plan.copy() - - for key in generic_keys: - param_name = key.rsplit(".", 1)[0] if "." in key else key - generic_param_name = re.sub(r"\d+", "*", param_name) - - if generic_param_name in tp_plan: - unused_rules.pop(generic_param_name, None) - unsharded_layers.discard(key) - elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: - unused_rules.pop(parent_param_name, None) - unsharded_layers.discard(key) - - if len(unused_rules) > 0: - logger.warning(f"The following TP rules were not applied on any of the layers: {unused_rules}") - if len(unsharded_layers) > 0: - logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") - - -def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size): - """Distribute a model according to the TP plan.""" - model._tp_size = tp_size - model._device_mesh = device_mesh - if distributed_config is not None: - if isinstance(distributed_config, dict): - distributed_config = DistributedConfig.from_dict(distributed_config) - model.config.distributed_config = distributed_config - # Set the new requested tp_plan on the model - if isinstance(tp_plan, dict): - model.tp_plan = tp_plan - model_plan = model.tp_plan - if model_plan is not None and _torch_distributed_available: - for v in model_plan.values(): - if v not in ALL_PARALLEL_STYLES: - raise ValueError(f"Unsupported tensor parallel style {v}. Supported styles are {ALL_PARALLEL_STYLES}") - for name, module in model.named_modules(): - if not getattr(module, "_is_hooked", False): - plan = _get_parameter_tp_plan(parameter_name=name, tp_plan=model_plan, is_weight=False) - add_tensor_parallel_hooks_to_module( - model, - module, - plan, - name, - device_mesh, - ) - module._is_hooked = True - return model diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 05708d400056..396480236a3a 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -39,6 +39,7 @@ from safetensors.torch import load as _safe_load_bytes from safetensors.torch import save_file as safe_save_file from torch import Tensor, nn +from torch.distributed.tensor import DTensor from torch.distributions import constraints from torch.utils.checkpoint import checkpoint @@ -52,9 +53,22 @@ revert_weight_conversion, ) from .distributed import DistributedConfig +from .distributed.fsdp import is_fsdp_enabled +from .distributed.sharding_utils import _dtensor_from_local_like +from .distributed.tensor_parallel import ( + _get_parameter_tp_plan, + verify_tp_plan, +) +from .distributed.utils import ( + _distributed_barrier, + distribute_model, + gather_full_state_dict, + init_device_mesh, + save_model_checkpoint_distributed, +) from .dynamic_module_utils import custom_object_save from .generation import CompileConfig, GenerationConfig -from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled, is_fsdp_enabled +from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled from .integrations.accelerate import ( _get_device_map, accelerate_disk_offload, @@ -75,15 +89,6 @@ from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward from .integrations.sdpa_paged import sdpa_attention_paged_forward -from .integrations.tensor_parallel import ( - ALL_PARALLEL_STYLES, - _get_parameter_tp_plan, - distribute_model, - gather_state_dict_for_save, - initialize_tensor_parallelism, - shard_and_distribute_module, - verify_tp_plan, -) from .loss.loss_utils import LOSS_MAPPING from .modeling_flash_attention_utils import ( FLASH_ATTENTION_COMPATIBILITY_MATRIX, @@ -93,7 +98,7 @@ ) from .modeling_rope_utils import ROPE_INIT_FUNCTIONS from .monkey_patching import apply_patches, patch_output_recorders -from .pytorch_utils import id_tensor_storage +from .pytorch_utils import _torch_distributed_available, id_tensor_storage from .quantizers import HfQuantizer from .quantizers.auto import get_hf_quantizer from .quantizers.quantizers_utils import get_module_from_name @@ -143,8 +148,6 @@ from ._typing import DeviceMeshLike -_torch_distributed_available = torch.distributed.is_available() - if is_sagemaker_mp_enabled(): import smdistributed.modelparallel.torch as smp from smdistributed.modelparallel import __version__ as SMP_VERSION @@ -182,6 +185,7 @@ class LoadStateDictConfig: dtype_plan: dict = field(default_factory=dict) hf_quantizer: HfQuantizer | None = None device_mesh: "DeviceMeshLike | None" = None + tp_plan: dict[str, str] | None = None weights_only: bool = True weight_mapping: list[WeightConverter | WeightRenaming] | None = None disable_mmap: bool | None = None @@ -1244,6 +1248,10 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH # For top-level models, this attribute is currently defined in respective model code. For base models, this attribute comes # from `config.base_model_tp_plan` during `post_init`. _tp_plan: dict[str, str] = None + # Sequence-parallel plan used when `distributed_config.enable_sequence_parallel` is set. For top-level models, this + # attribute is defined on the head class (e.g. `*ForCausalLM`). For base models, it comes from + # `config.base_model_sp_plan` during `post_init`. + _sp_plan: dict[str, str] = None # Tensor parallel degree to which model is sharded to _tp_size = None # A pipeline parallel plan specifying the layers which may not be present on all ranks when PP is enabled. For top-level @@ -1383,13 +1391,26 @@ def post_init(self): correctly in the case of composite models (that is, the top level model should know about those properties from its children). """ # Attach the different parallel plans and tied weight keys to the top-most model, so that everything is - # easily available - self._tp_plan, self._ep_plan, self._pp_plan = {}, {}, {} - # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config + # easily available. Seed with the class-level plans (e.g. `*ForCausalLM._tp_plan = {"lm_head": ...}`) + # before the instance attribute shadows them — task-head plans live on the head class. + cls_tp_plan = getattr(self, "_tp_plan", None) or {} + cls_sp_plan = getattr(self, "_sp_plan", None) or {} + cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} + self._tp_plan = dict(cls_tp_plan) + self._sp_plan = dict(cls_sp_plan) + self._ep_plan = {} + self._pp_plan = {} + self._fsdp_plan = dict(cls_fsdp_plan) + # If current model is a base model, attach `base_model_*_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} + self._sp_plan = self.config.base_model_sp_plan.copy() if self.config.base_model_sp_plan is not None else {} self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {} + self._fsdp_plan = ( + self.config.base_model_fsdp_plan.copy() if self.config.base_model_fsdp_plan is not None else {} + ) + # Current submodel should register its tied weights self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=False) # Current submodel should register its `_keep_in_fp32_modules` @@ -1411,8 +1432,12 @@ def post_init(self): self._ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_tp_plan", None): self._tp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_sp_plan", None): + self._sp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_pp_plan", None): self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_fsdp_plan", None): + self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) # Always attach the keys of the children (if the children's config says to NOT tie, then it's empty) if tied_keys := getattr(module, "all_tied_weights_keys", None): self.all_tied_weights_keys.update({f"{name}.{k}": f"{name}.{v}" for k, v in tied_keys.copy().items()}) @@ -1461,14 +1486,6 @@ def tp_plan(self, plan: dict[str, str] | None): if not isinstance(plan, dict): raise ValueError("Can only set a dictionary as `tp_plan`") - # Ensure the styles are all valid - for layer_pattern, parallel_style in plan.items(): - if parallel_style not in ALL_PARALLEL_STYLES: - raise ValueError( - f"Unsupported tensor parallel style '{parallel_style}' for layer '{layer_pattern}'. " - f"Supported styles are {list(ALL_PARALLEL_STYLES.keys())}" - ) - # Validate that the layer patterns match existing model structure. We check this by getting all parameter # names and seeing if any match the patterns model_param_names = [name for name, _ in self.named_parameters()] @@ -2065,6 +2082,8 @@ def _can_set_attn_implementation(cls) -> bool: """Detect whether the class supports setting its attention implementation dynamically. It is an ugly check based on opening the file, but avoids maintaining yet another property flag. """ + # Skip dynamic wrappers like FSDP2's FSDP, whose __module__ is inside torch.* + cls = next((k for k in cls.__mro__ if not k.__module__.startswith("torch.")), cls) class_module = sys.modules.get(cls.__module__) # Missing module entry (e.g. cleared by a test) or custom model in a jupyter notebook / repl -> do not allow to set it if class_module is None or not hasattr(class_module, "__file__"): @@ -2084,6 +2103,8 @@ def _can_set_experts_implementation(cls) -> bool: """Detect whether the class supports setting its experts implementation dynamically. It is an ugly check based on opening the file, but avoids maintaining yet another property flag. """ + # Skip dynamic wrappers like FSDP2's FSDP, whose __module__ is inside torch.* + cls = next((k for k in cls.__mro__ if not k.__module__.startswith("torch.")), cls) class_module = sys.modules.get(cls.__module__) # Missing module entry (e.g. cleared by a test) or custom model in a jupyter notebook / repl -> do not allow to set it if class_module is None or not hasattr(class_module, "__file__"): @@ -3304,6 +3325,7 @@ def save_pretrained( token: str | bool | None = None, save_peft_format: bool = True, save_original_format: bool = True, + distributed_checkpoint: bool = False, **kwargs, ): """ @@ -3369,7 +3391,7 @@ def save_pretrained( ) # we need to check against tp_size, not tp_plan, as tp_plan is substituted to the class one - if self._tp_size is not None and not is_huggingface_hub_greater_or_equal("0.31.4"): + if self.tp_size is not None and not is_huggingface_hub_greater_or_equal("0.31.4"): raise ImportError( "Saving a model with tensor parallelism requires `huggingface_hub` version 0.31.4 or higher." ) @@ -3379,9 +3401,13 @@ def save_pretrained( return os.makedirs(save_directory, exist_ok=True) + save_on_this_rank = is_main_process + if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is not None: + # DTensor gather materializes a full plain state dict only on global rank 0. + save_on_this_rank = save_on_this_rank and torch.distributed.get_rank() == 0 save_directory_path = os.fspath(save_directory) - if push_to_hub: + if push_to_hub and save_on_this_rank: commit_message = kwargs.pop("commit_message", None) repo_id = kwargs.pop("repo_id", save_directory_path.split(os.path.sep)[-1]) create_pr = kwargs.pop("create_pr", False) @@ -3406,46 +3432,72 @@ def save_pretrained( # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be # loaded from the Hub. - if self._auto_class is not None: + if self._auto_class is not None and save_on_this_rank: custom_object_save(self, save_directory, config=self.config) - # Save the config - if is_main_process: - if not _hf_peft_config_loaded: - model_to_save.config.save_pretrained(save_directory) - if self.can_generate(): - model_to_save.generation_config.save_pretrained(save_directory) + # Don't persist distributed_config in saved config — it's runtime-only + # (otherwise AutoConfig absorbs it on reload, preventing from_pretrained from seeing it as a kwarg). + # Keep a runtime copy around because TP/FSDP save helpers still rely on it after config serialization. + distributed_config = getattr(model_to_save.config, "distributed_config", None) + if distributed_config is not None: + del model_to_save.config.distributed_config - if _hf_peft_config_loaded: - logger.info( - "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved." - ) - state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict) + # Save the config + try: + if save_on_this_rank: + if not _hf_peft_config_loaded: + model_to_save.config.save_pretrained(save_directory) + if self.can_generate(): + model_to_save.generation_config.save_pretrained(save_directory) - if save_peft_format: + if _hf_peft_config_loaded: logger.info( - "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`." + "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved." ) - peft_state_dict = {} - for key, value in state_dict.items(): - peft_state_dict[f"base_model.model.{key}"] = value - state_dict = peft_state_dict + state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict) + + if save_peft_format: + logger.info( + "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`." + ) + peft_state_dict = {} + for key, value in state_dict.items(): + peft_state_dict[f"base_model.model.{key}"] = value + state_dict = peft_state_dict - active_adapter = self.active_adapters() + active_adapter = self.active_adapters() - if len(active_adapter) > 1: - raise ValueError( - "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one " - "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`" - ) - active_adapter = active_adapter[0] + if len(active_adapter) > 1: + raise ValueError( + "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one " + "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`" + ) + active_adapter = active_adapter[0] + + current_peft_config = self.peft_config[active_adapter] + current_peft_config.save_pretrained(save_directory) + finally: + if distributed_config is not None: + model_to_save.config.distributed_config = distributed_config - current_peft_config = self.peft_config[active_adapter] - current_peft_config.save_pretrained(save_directory) + if distributed_checkpoint: + if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is None: + raise ValueError( + "save_pretrained(distributed_checkpoint=True) requires the model to have been " + "initialized with a distributed_config (device_mesh is None)." + ) + save_model_checkpoint_distributed(self, save_directory) + return - # Get the model state_dict + # Get the model state_dict (handles FSDP unshard + TP gather in one call) + used_distributed_gather = False if state_dict is None: - state_dict = model_to_save.state_dict() + if getattr(self, "device_mesh", None) is not None: + # Pass self (not model_to_save) so device_mesh/tp_size/tp_plan are available + state_dict = gather_full_state_dict(self) + used_distributed_gather = True + else: + state_dict = model_to_save.state_dict() # if any model parameters are offloaded, we need to know it for later is_offloaded = False @@ -3471,10 +3523,6 @@ def save_pretrained( if ignore_key in state_dict: del state_dict[ignore_key] - # If model was sharded with TP, gather full tensors for saving - if self._tp_size is not None: - state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size) - # Remove tied weights as safetensors do not handle them state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save) @@ -3516,54 +3564,55 @@ def save_pretrained( filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and filename not in state_dict_split.filename_to_tensors - and is_main_process + and save_on_this_rank and reg.fullmatch(filename_no_suffix) is not None ): os.remove(full_filename) # Save the model - for shard_file, tensor_names in logging.tqdm( - state_dict_split.filename_to_tensors.items(), desc="Writing model shards" - ): - filename = os.path.join(save_directory, shard_file) - shard_state_dict = {} - for tensor_name in tensor_names: - # Get the tensor, and remove it from state_dict to avoid keeping the ref - tensor = state_dict.pop(tensor_name) - - # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern, - # but it would otherwise not be contained in the saved shard if we were to simply move the file - # or something - if is_offloaded and tensor.device.type == "meta": - tensor = load_offloaded_parameter(model_to_save, tensor_name) - - # only do contiguous after it's permuted correctly in case of TP - shard_state_dict[tensor_name] = tensor.contiguous() - - # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL, - # so it's not possible for now.... - # Write the shard to disk - safe_save_file(shard_state_dict, filename, metadata=metadata) - # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM) - del shard_state_dict - - if index is None: - path_to_weights = os.path.join(save_directory, weights_name) - logger.info(f"Model weights saved in {path_to_weights}") - else: - save_index_file = SAFE_WEIGHTS_INDEX_NAME - save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) - # Save the index as well - with open(save_index_file, "w", encoding="utf-8") as f: - content = json.dumps(index, indent=2, sort_keys=True) + "\n" - f.write(content) - logger.info( - f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " - f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " - f"index located at {save_index_file}." - ) + if save_on_this_rank: + for shard_file, tensor_names in logging.tqdm( + state_dict_split.filename_to_tensors.items(), desc="Writing model shards" + ): + filename = os.path.join(save_directory, shard_file) + shard_state_dict = {} + for tensor_name in tensor_names: + # Get the tensor, and remove it from state_dict to avoid keeping the ref + tensor = state_dict.pop(tensor_name) + + # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern, + # but it would otherwise not be contained in the saved shard if we were to simply move the file + # or something + if is_offloaded and tensor.device.type == "meta": + tensor = load_offloaded_parameter(model_to_save, tensor_name) + + # only do contiguous after it's permuted correctly in case of TP + shard_state_dict[tensor_name] = tensor.contiguous() + + # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL, + # so it's not possible for now.... + # Write the shard to disk + safe_save_file(shard_state_dict, filename, metadata=metadata) + # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM) + del shard_state_dict + + if index is None: + path_to_weights = os.path.join(save_directory, weights_name) + logger.info(f"Model weights saved in {path_to_weights}") + else: + save_index_file = SAFE_WEIGHTS_INDEX_NAME + save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) - if push_to_hub: + if push_to_hub and save_on_this_rank: # Eventually create an empty model card model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token) @@ -3579,6 +3628,13 @@ def save_pretrained( create_pr=create_pr, ) + # `gather_full_state_dict` concentrates the full state on rank 0 only; + # other ranks then loop over an empty shard list and would race ahead + # of rank 0's safetensors writes. Barrier so any subsequent + # `from_pretrained` on this path sees the consolidated files. + if used_distributed_gather: + _distributed_barrier() + @wraps(PushToHubMixin.push_to_hub) def push_to_hub(self, *args, **kwargs): tags = self.model_tags if self.model_tags is not None else [] @@ -3972,13 +4028,22 @@ def from_pretrained( max_memory (`Dict`, *optional*): A dictionary device identifier to maximum memory if using `device_map`. Will default to the maximum memory available for each GPU and the available CPU RAM if unset. - tp_plan (`Optional[Union[dict, str]]`, *optional*): - A torch tensor parallel plan, see [here](https://pytorch.org/tutorials/intermediate/TP_tutorial.html). Use `tp_plan="auto"` to - use the predefined plan based on the model. If it's a dict, then it should match between module names and desired layout. - Note that if you use it, you should launch your script accordingly with `torchrun [args] script.py`. This will be much - faster than using a `device_map`, but has limitations. - tp_size (`str`, *optional*): - A torch tensor parallel degree. If not provided would default to world size. + distributed_config ([`DistributedConfig`], *optional*): + Configuration for native distributed training (FSDP2 + TP) via `torch.distributed`. Mutually + exclusive with `quantization_config` (for now) and `device_map`. When set, accelerate is not used for + device placement or dispatch. Launch with `torchrun --nproc_per_node=N script.py`. + + Accepts `tp_size`, `tp_plan`, `fsdp_size`, `fsdp_plan`. When a size is specified without a + plan, the plan defaults to the model's predefined behavior: TP uses the model's predefined + tensor parallel sharding plan, FSDP wraps each transformer layer individually with FSDP2 + (`fully_shard`). Both plans also accept a `dict` for manual control: `tp_plan` maps + parameter names to parallel styles (e.g. `{"model.layers.*.self_attn.q_proj": "colwise"}`), + `fsdp_plan` maps module names to wrap (e.g. `{"model.layers.0": {}, "model.layers.1": {}}`). + + Examples: + - TP-only: `DistributedConfig(tp_size=4)` + - FSDP-only: `DistributedConfig(fsdp_size=4)` + - 2D parallel: `DistributedConfig(tp_size=2, fsdp_size=2)` on 4 GPUs device_mesh (`torch.distributed.DeviceMesh`, *optional*): A torch device mesh. If not provided would default to world size. Used only for tensor parallel for now. If provided, it has to contain dimension named `"tp"` in case it's > 1 dimensional, this dimension will be used for tensor parallelism @@ -4072,8 +4137,6 @@ def from_pretrained( adapter_name = kwargs.pop("adapter_name", "default") generation_config = kwargs.pop("generation_config", None) gguf_file = kwargs.pop("gguf_file", None) - tp_plan = kwargs.pop("tp_plan", None) - tp_size = kwargs.pop("tp_size", None) distributed_config: DistributedConfig = kwargs.pop("distributed_config", None) device_mesh = kwargs.pop("device_mesh", None) trust_remote_code = kwargs.pop("trust_remote_code", None) @@ -4082,8 +4145,17 @@ def from_pretrained( kernel_config = kwargs.pop("kernel_config", None) key_mapping = kwargs.pop("key_mapping", None) - if distributed_config is not None and tp_plan is None: - tp_plan = "auto" + if distributed_config is not None: + if device_map is not None: + raise ValueError( + "`distributed_config` and `device_map` are mutually exclusive. " + "`distributed_config` handles device placement natively via torch.distributed." + ) + # NOTE(3outeille): support quantization (fp4/fp8) with distributed training later + if quantization_config is not None: + raise ValueError( + "Quantization is not currently supported with distributed training. Please disable quantization or distributed_config." + ) # Not used anymore -- remove them from the kwargs for name in ["mirror", "_fast_init", "low_cpu_mem_usage", "from_tf", "from_flax", "offload_state_dict"]: @@ -4114,17 +4186,24 @@ def from_pretrained( "`state_dict` cannot be passed together with a model name or a `gguf_file`. Use one of the two loading strategies." ) - if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")): - logger.info( - "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. " - "If your plan is to load the model on each device, you should set device_map={" - ": PartialState().process_index} where PartialState comes from accelerate library" - ) + if distributed_config is not None: + device_mesh = init_device_mesh(distributed_config) + # When using any distributed setup, we simply set the device_map here to the current rank so that we correctly + # load the params on the right device + if device_mesh.device_type == "cpu": + device_map = {"": torch.device("cpu")} + else: + device_map = {"": torch.device(device_mesh.device_type, int(os.environ["LOCAL_RANK"]))} + else: + # Accelerate path + if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")): + logger.info( + "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. " + "If your plan is to load the model on each device, you should set device_map={" + ": PartialState().process_index} where PartialState comes from accelerate library" + ) - if tp_plan is not None or tp_size is not None: # TP warnings, and setup - device_map, device_mesh, tp_size = initialize_tensor_parallelism( - tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map - ) + device_map = check_and_set_device_map(device_map) # validate & normalize (requires accelerate) if gguf_file is not None and not is_accelerate_available(): raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.") @@ -4137,7 +4216,6 @@ def from_pretrained( download_kwargs_with_commit, **adapter_kwargs, ) - device_map = check_and_set_device_map(device_map) # warn, error and fix the device map user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class} if from_pipeline is not None: @@ -4265,17 +4343,24 @@ def from_pretrained( # instantiated model, as the flags can be modified by instances sometimes) dtype_plan = model._get_dtype_plan(dtype) - # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively + # Obtain the weight conversion mapping for this model if any are registered weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) - if _torch_distributed_available and device_mesh is not None: # add hooks to nn.Modules: no weights - model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size) - - # Prepare the full device map - if device_map is not None: + if distributed_config is not None: + # Tie weights before sharding so `apply_fully_shard_data_parallel` / + # `apply_tensor_parallel` see the shared-parameter graph and can route tied + # entries (e.g. `lm_head` -> `embed_tokens`) correctly. `_finalize_model_loading` + # re-runs `tie_weights` after the checkpoint is loaded to handle missing-key edge cases. + model.tie_weights() + model = distribute_model(model, distributed_config, device_mesh) + elif device_map is not None: + # Expand device_map if it was passed as a `str`, i.e. `device_map="auto"` device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) # Finalize model weight initialization + active_tp_plan = ( + getattr(model, "_tp_plan", None) if getattr(distributed_config, "tp_size", None) is not None else None + ) load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, @@ -4287,6 +4372,7 @@ def from_pretrained( dtype_plan=dtype_plan, hf_quantizer=hf_quantizer, device_mesh=device_mesh, + tp_plan=active_tp_plan, weights_only=weights_only, weight_mapping=weight_conversions, use_safetensors=use_safetensors, @@ -4355,7 +4441,7 @@ def _load_pretrained_model( expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys if logger.level >= logging.WARNING: - verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None)) + verify_tp_plan(expected_keys, load_config.tp_plan) # This offload index if for params explicitly on the "disk" in the device_map disk_offload_index = None @@ -4427,7 +4513,7 @@ def _load_pretrained_model( model=model, state_dict=merged_state_dict, load_config=load_config, - tp_plan=model.tp_plan, + tp_plan=load_config.tp_plan, disk_offload_index=disk_offload_index, ) @@ -4579,8 +4665,8 @@ def tp_size(self): """ Returns the model's tensor parallelism degree. """ - # if None, the model didn't undergo tensor parallel sharding - return self._tp_size + dc = getattr(self.config, "distributed_config", None) + return dc.tp_size if dc is not None else None @property def supports_pp_plan(self): @@ -4735,15 +4821,20 @@ def _move_missing_keys_from_meta_to_device( # will be re-initialized for nothing (which can be quite long) for key in missing_keys - self.all_tied_weights_keys.keys(): param = self.get_parameter_or_buffer(key) - param_device = get_device(device_map, key, valid_torch_device=True) - value = torch.empty_like(param, device=param_device) - # For TP, we may need to shard the param - if device_mesh is not None: - shard_and_distribute_module( - self, value, param, key, None, False, device_mesh.get_local_rank(), device_mesh + if isinstance(param, DTensor): + # DTensor from parallelize_module on meta — materialize on actual device + local_value = torch.empty( + param._local_tensor.shape, + dtype=param.dtype, + device=torch.device(param.device_mesh.device_type, torch.cuda.current_device()), ) - # Otherwise, just move it to device + new_dtensor = _dtensor_from_local_like(local_value, param) + with torch.no_grad(): + new_param = torch.nn.Parameter(new_dtensor, requires_grad=param.requires_grad) + torch.utils.swap_tensors(param, new_param) else: + param_device = get_device(device_map, key, valid_torch_device=True) + value = torch.empty_like(param, device=param_device) _load_parameter_into_model(self, key, value) # We need to move back non-persistent buffers as well, as they are not part of loaded weights anyway for key, buffer in self.named_non_persistent_buffers(): diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 421119b33deb..1fb30e564e1c 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -615,8 +615,10 @@ def forward( @auto_docstring class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index f3ff9f15b103..0875127a3215 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -395,7 +395,9 @@ def forward( class AfmoeForCausalLM(LlamaForCausalLM, AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/apertus/configuration_apertus.py b/src/transformers/models/apertus/configuration_apertus.py index 1e0122160b30..c7fcc4159c06 100644 --- a/src/transformers/models/apertus/configuration_apertus.py +++ b/src/transformers/models/apertus/configuration_apertus.py @@ -49,11 +49,26 @@ class ApertusConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class ApertusConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 131072 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 7d14dd3d14c8..8c5757377b03 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -422,8 +422,10 @@ def forward( @auto_docstring class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/apertus/modular_apertus.py b/src/transformers/models/apertus/modular_apertus.py index 3c9eb6d8b6ea..075a6de4776b 100644 --- a/src/transformers/models/apertus/modular_apertus.py +++ b/src/transformers/models/apertus/modular_apertus.py @@ -67,11 +67,26 @@ class ApertusConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -79,6 +94,12 @@ class ApertusConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 131072 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/arcee/configuration_arcee.py b/src/transformers/models/arcee/configuration_arcee.py index ee711d608204..e1d957bcac43 100644 --- a/src/transformers/models/arcee/configuration_arcee.py +++ b/src/transformers/models/arcee/configuration_arcee.py @@ -49,9 +49,23 @@ class ArceeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -59,6 +73,12 @@ class ArceeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 2560 intermediate_size: int = 18432 diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 8d2d05bf2952..05ebab40a879 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -424,8 +424,10 @@ def forward( @auto_docstring(checkpoint="arcee-ai/AFM-4.5B") class ArceeForCausalLM(ArceePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/arcee/modular_arcee.py b/src/transformers/models/arcee/modular_arcee.py index 91cd8e13f1ed..a382b2e7191b 100644 --- a/src/transformers/models/arcee/modular_arcee.py +++ b/src/transformers/models/arcee/modular_arcee.py @@ -53,9 +53,23 @@ class ArceeConfig(LlamaConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } vocab_size: int = 32000 diff --git a/src/transformers/models/aria/configuration_aria.py b/src/transformers/models/aria/configuration_aria.py index 7694656905dd..dfc1b8e74258 100644 --- a/src/transformers/models/aria/configuration_aria.py +++ b/src/transformers/models/aria/configuration_aria.py @@ -44,10 +44,25 @@ class AriaTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -55,6 +70,12 @@ class AriaTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index e4fab906e561..91fcf5d106da 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -758,8 +758,10 @@ def forward( @auto_docstring class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: AriaTextConfig): super().__init__(config) diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py index e6e7c5abd3c0..9b3c978b756c 100644 --- a/src/transformers/models/aria/modular_aria.py +++ b/src/transformers/models/aria/modular_aria.py @@ -111,10 +111,10 @@ class AriaTextConfig(LlamaConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } intermediate_size: int = 4096 diff --git a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py index 8eff94824ed6..26972fcf18b2 100644 --- a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py @@ -447,6 +447,7 @@ def forward(self, audio_features): ) class AudioFlamingo3Model(AudioFlamingo3PreTrainedModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None diff --git a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py index 9eecd9d202cd..86be1ffa376f 100644 --- a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py @@ -182,6 +182,7 @@ def __init__(self, config: AudioFlamingo3Config): ) class AudioFlamingo3Model(VoxtralModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index 2b42a603d841..c3b3a8a24bf1 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1071,8 +1071,10 @@ def _update_mamba_mask(self, attention_mask, past_key_values): @auto_docstring class BambaForCausalLM(BambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/bitnet/modeling_bitnet.py b/src/transformers/models/bitnet/modeling_bitnet.py index 14c1581b250f..4994407d74c1 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -423,7 +423,9 @@ def forward( class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None _pp_plan = None + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/bitnet/modular_bitnet.py b/src/transformers/models/bitnet/modular_bitnet.py index 7f3a248766b2..c8376f131859 100644 --- a/src/transformers/models/bitnet/modular_bitnet.py +++ b/src/transformers/models/bitnet/modular_bitnet.py @@ -110,6 +110,7 @@ class BitNetModel(LlamaModel): class BitNetForCausalLM(LlamaForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None + _sp_plan = None _pp_plan = None def forward( diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index d52a3e008427..4c05dce75e47 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -53,16 +53,39 @@ class CohereConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index b8bf50af9bf4..309ddc142dc3 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -454,8 +454,10 @@ def forward( @auto_docstring class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/cohere2/configuration_cohere2.py b/src/transformers/models/cohere2/configuration_cohere2.py index 48c2df360354..19197e26204d 100644 --- a/src/transformers/models/cohere2/configuration_cohere2.py +++ b/src/transformers/models/cohere2/configuration_cohere2.py @@ -52,10 +52,25 @@ class Cohere2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -63,6 +78,12 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index 61bac127e16f..ce14624cb982 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -432,8 +432,10 @@ def forward( @auto_docstring class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/cohere2/modular_cohere2.py b/src/transformers/models/cohere2/modular_cohere2.py index b40a5056232b..ecd41c3e86d4 100644 --- a/src/transformers/models/cohere2/modular_cohere2.py +++ b/src/transformers/models/cohere2/modular_cohere2.py @@ -73,10 +73,25 @@ class Cohere2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -84,6 +99,12 @@ class Cohere2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 8192 intermediate_size: int = 22528 diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py index 9c1abfc06cd7..f48a04a5939f 100644 --- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py @@ -74,16 +74,36 @@ class Cohere2MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + # Depending on layers, `mlp` can be a MoE or an MLP layer + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + # TP + Sequence Parallelism plan (for training). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + # Depending on layers, `mlp` can be a MoE or an MLP layer - so they are updated by `_update_sp_plan` to avoid collisions + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan (see Qwen3Config.base_model_fsdp_plan for shape rationale). + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } vocab_size: int = 256000 hidden_size: int = 8192 @@ -142,7 +162,6 @@ def __post_init__(self, **kwargs): for i in range(self.num_hidden_layers - first_k_dense_replace) ] self.layer_types = prefix_layers + rest_layers - self.validate_layer_type() if self.mlp_layer_types is None: @@ -151,6 +170,30 @@ def __post_init__(self, **kwargs): ] super().__post_init__(**kwargs) + self._update_sp_plan() + + def _update_sp_plan(self): + """Depending on layers, `mlp` can be a MoE or an MLP layer - so we update the plan dynamically here to avoid collisions""" + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + # This is a MLP layer + if mlp_type == "dense": + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + # This is a MoE layer + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) __all__ = ["Cohere2MoeConfig"] diff --git a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py index 1c9aaa5b42ca..d25fc140835e 100644 --- a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py @@ -581,8 +581,10 @@ def forward( @auto_docstring class Cohere2MoeForCausalLM(Cohere2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index dd0a1f0aa749..99f643e93b44 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -549,7 +549,9 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(CsmPreTrainedModel, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None _pp_plan = None + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/csm/modular_csm.py b/src/transformers/models/csm/modular_csm.py index c82636613ee1..f1e95384ab52 100644 --- a/src/transformers/models/csm/modular_csm.py +++ b/src/transformers/models/csm/modular_csm.py @@ -274,6 +274,7 @@ def forward(self, hidden_states, codebook_indices=None): class CsmDepthDecoderForCausalLM(LlamaForCausalLM, GenerationMixin): _tied_weights_keys = None _tp_plan = None + _sp_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/cwm/configuration_cwm.py b/src/transformers/models/cwm/configuration_cwm.py index ecc3743da19d..6eb7175ad020 100644 --- a/src/transformers/models/cwm/configuration_cwm.py +++ b/src/transformers/models/cwm/configuration_cwm.py @@ -49,10 +49,25 @@ class CwmConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,6 +75,12 @@ class CwmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 6144 intermediate_size: int = 21504 diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index 3e0eb0504be0..b20fd9b7436f 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -426,8 +426,10 @@ def forward( @auto_docstring class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/data2vec/modeling_data2vec_audio.py b/src/transformers/models/data2vec/modeling_data2vec_audio.py index bb3af4ecf25a..048747656b9e 100755 --- a/src/transformers/models/data2vec/modeling_data2vec_audio.py +++ b/src/transformers/models/data2vec/modeling_data2vec_audio.py @@ -29,8 +29,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index 58735fb55c0b..a984b4585a55 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -642,7 +642,9 @@ def load_balancing_loss_func( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/dbrx/modular_dbrx.py b/src/transformers/models/dbrx/modular_dbrx.py index d737d59e1a8b..548af86ef4f6 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -430,7 +430,9 @@ def forward( class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 1b8005f8efef..49b3cb9c6387 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -54,18 +54,39 @@ class DeepseekV2Config(PreTrainedConfig): base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + # MLA: don't shard q_a_proj / kv_a_proj_with_mqa — their outputs feed + # layernorms whose weights are full-size. Only b-projections (after + # the norm) and o_proj are sharded. + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + # Shared experts output must stay Replicate to match experts output + # (they're summed inside the MoE block, before the outer allgather_split + # handles the SP boundary). + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +94,12 @@ class DeepseekV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index 3ef8266218f7..1f9ba33ec6a5 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -541,8 +541,10 @@ def forward( @auto_docstring class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 5644c7dc2990..a6c6494c0415 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -69,18 +69,39 @@ class DeepseekV2Config(LlamaConfig): base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + # MLA: don't shard q_a_proj / kv_a_proj_with_mqa — their outputs feed + # layernorms whose weights are full-size. Only b-projections (after + # the norm) and o_proj are sharded. + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_b_proj": "colwise", + "layers.*.self_attn.kv_b_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + # Shared experts output must stay Replicate to match experts output + # (they're summed inside the MoE block, before the outer allgather_split + # handles the SP boundary). + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } model_type = "deepseek_v2" diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 4178547a5ff2..32e67ab4d90a 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -49,21 +49,25 @@ class DeepseekV3Config(PreTrainedConfig): model_type = "deepseek_v3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", } @@ -111,5 +115,20 @@ def __post_init__(self, **kwargs): self.head_dim = self.qk_rope_head_dim super().__post_init__(**kwargs) + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs + __all__ = ["DeepseekV3Config"] diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index fe3acd9aeddd..8d828320b19b 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -634,8 +634,10 @@ def forward( @auto_docstring class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 8f8d818d8ced..9475e8d6f1d0 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py @@ -125,7 +125,13 @@ class DeepseekV4Config(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } vocab_size: int = 129280 diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index d381ed930458..46a4d965ba55 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -1398,8 +1398,10 @@ def load_balancing_loss_func( @auto_docstring class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/dia/generation_dia.py b/src/transformers/models/dia/generation_dia.py index d22b2fff0d8d..76506e66f00c 100644 --- a/src/transformers/models/dia/generation_dia.py +++ b/src/transformers/models/dia/generation_dia.py @@ -18,6 +18,7 @@ import torch import torch.distributed as dist +from ...distributed.fsdp import is_fsdp_managed_module from ...generation.logits_process import ( DiaClassifierFreeGuidanceLogitsProcessor, DiaEOSChannelFilterLogitsProcessor, @@ -29,7 +30,6 @@ from ...generation.streamers import BaseStreamer from ...generation.utils import GenerateOutput, GenerationConfig, GenerationMixin, GenerationMode from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_utils import PreTrainedModel from ...utils import logging diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index d80ccd572dc3..b822da50e2af 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -660,8 +660,10 @@ def forward( @auto_docstring class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index 8518d9021458..feac227b5a97 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -55,14 +55,14 @@ class DogeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.dt_proj": "rowwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.dt_proj": "rowwise_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.mlp.router_gate": "colwise_gather_output", - "layers.*.mlp.down_embed": "rowwise_split_input", - "layers.*.mlp.up_embed": "rowwise_split_input", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.router_gate": "colwise_allgather", + "layers.*.mlp.down_embed": "vocab_allreduce", + "layers.*.mlp.up_embed": "vocab_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -70,6 +70,12 @@ class DogeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32768 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 4aad59b52a9a..1f112465369d 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -716,8 +716,10 @@ def load_balancing_loss_func( @auto_docstring class DogeForCausalLM(DogePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index 8b78126c0a00..bc7e22089ea3 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -84,14 +84,14 @@ class DogeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.dt_proj": "rowwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.dt_proj": "rowwise_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.mlp.router_gate": "colwise_gather_output", - "layers.*.mlp.down_embed": "rowwise_split_input", - "layers.*.mlp.up_embed": "rowwise_split_input", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.router_gate": "colwise_allgather", + "layers.*.mlp.down_embed": "vocab_allreduce", + "layers.*.mlp.up_embed": "vocab_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -99,6 +99,12 @@ class DogeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32768 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py index 4d568bf4a565..9846f961250f 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -51,18 +51,14 @@ class Dots1Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { @@ -70,6 +66,13 @@ class Dots1Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 95b21258ffd5..05b1a6d20b4c 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -569,8 +569,10 @@ def forward( @auto_docstring class Dots1ForCausalLM(Dots1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index 06402d63e28c..94bd074871c0 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -65,18 +65,14 @@ class Dots1Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { @@ -84,6 +80,13 @@ class Dots1Config(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index eab9ee97da93..bd3f3011a539 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1275,8 +1275,10 @@ def forward( @auto_docstring class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Emu3TextConfig def __init__(self, config): diff --git a/src/transformers/models/ernie4_5/configuration_ernie4_5.py b/src/transformers/models/ernie4_5/configuration_ernie4_5.py index 896fb0b99402..5d805e1114e6 100644 --- a/src/transformers/models/ernie4_5/configuration_ernie4_5.py +++ b/src/transformers/models/ernie4_5/configuration_ernie4_5.py @@ -50,10 +50,25 @@ class Ernie4_5Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class Ernie4_5Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 hidden_size: int = 1024 intermediate_size: int = 3072 diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index ae533c7c6ef8..73f15f713be4 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -422,8 +422,10 @@ def forward( @auto_docstring class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py index 0c0c0edbb760..035daabd966d 100644 --- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py @@ -66,16 +66,14 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -83,6 +81,12 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 pad_token_id: int | None = 0 bos_token_id: int | None = 1 diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index 9c106a90010d..939201f4ff19 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -654,8 +654,10 @@ def load_balancing_loss_func( @auto_docstring class Ernie4_5_MoeForCausalLM(Ernie4_5_MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py index e4eea836f107..34caaa02dbeb 100644 --- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py @@ -50,9 +50,9 @@ class Ernie4_5_VLMoeVisionConfig(PreTrainedConfig): base_model_tp_plan = { "blocks.*.attn.qkv": "colwise", - "blocks.*.attn.proj": "rowwise", + "blocks.*.attn.proj": "rowwise_allreduce", "blocks.*.mlp.fc1": "colwise", - "blocks.*.mlp.fc2": "rowwise", + "blocks.*.mlp.fc2": "rowwise_allreduce", } intermediate_size: int = 4 * 1280 temporal_merge_size: int = 2 @@ -86,13 +86,13 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -100,6 +100,12 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 pad_token_id: int | None = None bos_token_id: int | None = None diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py index 3736acd51103..0ec36d704726 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py @@ -94,9 +94,9 @@ class Ernie4_5_VLMoeVisionConfig(Qwen2VLVisionConfig): base_model_tp_plan = { "blocks.*.attn.qkv": "colwise", - "blocks.*.attn.proj": "rowwise", + "blocks.*.attn.proj": "rowwise_allreduce", "blocks.*.mlp.fc1": "colwise", - "blocks.*.mlp.fc2": "rowwise", + "blocks.*.mlp.fc2": "rowwise_allreduce", } hidden_size: int = 1280 @@ -134,13 +134,13 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section"} diff --git a/src/transformers/models/eurobert/configuration_eurobert.py b/src/transformers/models/eurobert/configuration_eurobert.py index f64c4f7e5a11..ab2c33a28828 100644 --- a/src/transformers/models/eurobert/configuration_eurobert.py +++ b/src/transformers/models/eurobert/configuration_eurobert.py @@ -54,10 +54,25 @@ class EuroBertConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -65,6 +80,12 @@ class EuroBertConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 768 intermediate_size: int = 3072 diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index b93dd0649f14..655d10ab123d 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -408,7 +408,9 @@ def forward( @auto_docstring class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/eurobert/modular_eurobert.py b/src/transformers/models/eurobert/modular_eurobert.py index 19f965c84f73..b40c4d043707 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -141,7 +141,9 @@ def forward( @auto_docstring class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/exaone4/configuration_exaone4.py b/src/transformers/models/exaone4/configuration_exaone4.py index f29cab8dd8ea..fd3bade9f662 100644 --- a/src/transformers/models/exaone4/configuration_exaone4.py +++ b/src/transformers/models/exaone4/configuration_exaone4.py @@ -63,12 +63,27 @@ class Exaone4Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +91,12 @@ class Exaone4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index fab10b9b6937..d6aa12baedad 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -440,8 +440,10 @@ def forward( @auto_docstring class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/exaone4/modular_exaone4.py b/src/transformers/models/exaone4/modular_exaone4.py index c6d9202170a0..7a88b77629f1 100644 --- a/src/transformers/models/exaone4/modular_exaone4.py +++ b/src/transformers/models/exaone4/modular_exaone4.py @@ -92,12 +92,27 @@ class Exaone4Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -105,6 +120,12 @@ class Exaone4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 81ded9366cdb..83a873314700 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -69,19 +69,25 @@ class ExaoneMoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index a7f80fc979c4..d2d8543f1417 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -563,8 +563,10 @@ def forward( @auto_docstring class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/exaone_moe/modular_exaone_moe.py b/src/transformers/models/exaone_moe/modular_exaone_moe.py index 75ec2b0bfd27..8d50c3db2295 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -80,6 +80,14 @@ class ExaoneMoeConfig(Exaone4Config): >>> configuration = model.config ```""" + base_model_sp_plan = None + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 102400 hidden_size: int = 4096 intermediate_size: int = 16384 diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 3be9f882008e..2ea3c884437e 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1166,8 +1166,10 @@ def _update_mamba_mask(self, attention_mask, past_key_values): @auto_docstring class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index 7b08a79b801b..a89d9f0452e8 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -50,13 +50,10 @@ class FlexOlmoConfig(PreTrainedConfig): attribute_map = {"num_local_experts": "num_experts"} default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +61,12 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index 100e6fa35554..a9403a4110b0 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -597,8 +597,10 @@ def load_balancing_loss_func( @auto_docstring class FlexOlmoForCausalLM(FlexOlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index 01f32227f31f..c7b73dea8f15 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -60,13 +60,10 @@ class FlexOlmoConfig(PreTrainedConfig): attribute_map = {"num_local_experts": "num_experts"} default_theta = 500000.0 base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -74,6 +71,12 @@ class FlexOlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/gemma/configuration_gemma.py b/src/transformers/models/gemma/configuration_gemma.py index 0ef2e8b31b8d..7b80a3fda63e 100644 --- a/src/transformers/models/gemma/configuration_gemma.py +++ b/src/transformers/models/gemma/configuration_gemma.py @@ -50,10 +50,25 @@ class GemmaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class GemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 3072 intermediate_size: int = 24576 diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index c6c5a55b8790..f64702b55f19 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -450,8 +450,10 @@ def forward( @auto_docstring class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/gemma/modular_gemma.py b/src/transformers/models/gemma/modular_gemma.py index 25f436473fbe..8a24552eaa36 100644 --- a/src/transformers/models/gemma/modular_gemma.py +++ b/src/transformers/models/gemma/modular_gemma.py @@ -69,10 +69,25 @@ class GemmaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -80,6 +95,12 @@ class GemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 3072 intermediate_size: int = 24576 diff --git a/src/transformers/models/gemma2/configuration_gemma2.py b/src/transformers/models/gemma2/configuration_gemma2.py index 11d7b012099f..8ff4fc88629a 100644 --- a/src/transformers/models/gemma2/configuration_gemma2.py +++ b/src/transformers/models/gemma2/configuration_gemma2.py @@ -54,10 +54,25 @@ class Gemma2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -65,6 +80,12 @@ class Gemma2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 20673571b2d2..46eb89681619 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -476,8 +476,10 @@ def forward( @auto_docstring class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/gemma2/modular_gemma2.py b/src/transformers/models/gemma2/modular_gemma2.py index 2edd9ef5f101..d400331af615 100644 --- a/src/transformers/models/gemma2/modular_gemma2.py +++ b/src/transformers/models/gemma2/modular_gemma2.py @@ -81,10 +81,25 @@ class Gemma2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -92,6 +107,12 @@ class Gemma2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/gemma3/configuration_gemma3.py b/src/transformers/models/gemma3/configuration_gemma3.py index f25c9cef21fd..adf71ed56bba 100644 --- a/src/transformers/models/gemma3/configuration_gemma3.py +++ b/src/transformers/models/gemma3/configuration_gemma3.py @@ -61,12 +61,27 @@ class Gemma3TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -74,6 +89,12 @@ class Gemma3TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_208 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -99,6 +120,7 @@ class Gemma3TextConfig(PreTrainedConfig): final_logit_softcapping: float | None = None attn_logit_softcapping: float | None = None use_bidirectional_attention: bool | None = False + default_theta = {"global": 1_000_000.0, "local": 10_000.0} def __post_init__(self, **kwargs): diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 1a7b79afcac1..b175dc6b67a4 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -585,8 +585,10 @@ def forward( @auto_docstring class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma3TextConfig def __init__(self, config: Gemma3TextConfig): diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index 82afbc884c65..e83243f21c4a 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -87,13 +87,35 @@ class Gemma3TextConfig(Gemma2Config, PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size: int = 262_208 diff --git a/src/transformers/models/gemma3n/configuration_gemma3n.py b/src/transformers/models/gemma3n/configuration_gemma3n.py index e61c5f0038e7..973ff753c28b 100644 --- a/src/transformers/models/gemma3n/configuration_gemma3n.py +++ b/src/transformers/models/gemma3n/configuration_gemma3n.py @@ -78,23 +78,27 @@ class Gemma3nTextConfig(PreTrainedConfig): model_type = "gemma3n_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.v_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_400 hidden_size: int = 2048 intermediate_size: int | list[int] = 16_384 @@ -117,6 +121,7 @@ class Gemma3nTextConfig(PreTrainedConfig): sliding_window: int = 512 layer_types: list[str] | None = None final_logit_softcapping: float = 30.0 + default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size_per_layer_input: int = 262_144 hidden_size_per_layer_input: int = 256 diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index ff674c2eee72..dafdd8bd27a8 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1827,8 +1827,10 @@ def project_per_layer_inputs( @auto_docstring(custom_intro="The base Gemma 3n language model with a language modeling head.") class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma3nTextConfig def __init__(self, config: Gemma3nTextConfig): diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index 4ce215996fb4..2cbd72a61ed0 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -116,17 +116,22 @@ class Gemma3nTextConfig(Gemma3TextConfig): model_type = "gemma3n_text" base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.v_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + default_theta = {"global": 1_000_000.0, "local": 10_000.0} vocab_size: int = 262_400 diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index 66d1548fe453..c1c8cd1f26d4 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -125,28 +125,25 @@ class Gemma4TextConfig(PreTrainedConfig): model_type = "gemma4_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + # q/k use allgather because gemma4 has q_norm/k_norm with full-sized weights + # that can't match sharded q/k outputs. + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.experts.gate_up_proj": "packed_colwise", - "layers.*.experts.down_proj": "rowwise", - "layers.*.experts": "moe_tp_experts", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_ep_plan = { # EP plan for google/gemma-4-26B-A4B-it: do not tp in attention (num_global_key_value_heads=2 too small to partition) "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.router": "ep_router", "layers.*.experts.gate_up_proj": "grouped_gemm", "layers.*.experts.down_proj": "grouped_gemm", - "layers.*.experts": "moe_tp_experts", + "layers.*.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -154,6 +151,12 @@ class Gemma4TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_144 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -242,12 +245,10 @@ class Gemma4VisionConfig(PreTrainedConfig): "encoder.layers.*.self_attn.q_proj": "colwise", "encoder.layers.*.self_attn.k_proj": "colwise", "encoder.layers.*.self_attn.v_proj": "colwise", - "encoder.layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "encoder.layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "encoder.layers.*.self_attn.o_proj": "rowwise", + "encoder.layers.*.self_attn.o_proj": "rowwise_allreduce", "encoder.layers.*.mlp.gate_proj": "colwise", "encoder.layers.*.mlp.up_proj": "colwise", - "encoder.layers.*.mlp.down_proj": "rowwise", + "encoder.layers.*.mlp.down_proj": "rowwise_allreduce", } default_theta = 100.0 diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index 8cd171841fa5..42e34c48eb70 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1810,8 +1810,10 @@ def project_per_layer_inputs( @auto_docstring(custom_intro="The base Gemma 4 language model with a language modeling head.") class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma4TextConfig base_model_prefix = "model" diff --git a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py index 8b7b6b1a8b67..4c300686a0e7 100644 --- a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py +++ b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py @@ -110,6 +110,7 @@ def _init_weights(self, module): class Gemma4AssistantForCausalLM(Gemma4AssistantPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} + _fsdp_plan = {"lm_head": "keep_full_weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: Gemma4AssistantConfig): diff --git a/src/transformers/models/glm/configuration_glm.py b/src/transformers/models/glm/configuration_glm.py index 98525012a23b..67275679ab8b 100644 --- a/src/transformers/models/glm/configuration_glm.py +++ b/src/transformers/models/glm/configuration_glm.py @@ -43,9 +43,23 @@ class GlmConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "packed_colwise", # fused gate/up shards stay local for chunk + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "packed_colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -53,6 +67,12 @@ class GlmConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151552 hidden_size: int = 4096 intermediate_size: int = 13696 diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index 712202580943..461703ab81cd 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -439,8 +439,10 @@ def forward( @auto_docstring class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4/configuration_glm4.py b/src/transformers/models/glm4/configuration_glm4.py index f33129607fec..8b2755323f68 100644 --- a/src/transformers/models/glm4/configuration_glm4.py +++ b/src/transformers/models/glm4/configuration_glm4.py @@ -43,9 +43,23 @@ class Glm4Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": "vocab_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -53,6 +67,12 @@ class Glm4Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151552 hidden_size: int = 4096 intermediate_size: int = 13696 diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index e99930ae57f6..aa3d6e62a124 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -444,8 +444,10 @@ def forward( @auto_docstring class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py index a18123e90b33..bd94a492bb66 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -57,22 +57,27 @@ class Glm4MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", # NOTE(3outeille): This needs to be right after down_proj in the dict. Otherwise, the pattern model.layers.*.mlp.experts will have priority over model.layers.*.mlp.experts.down_proj which will assign a wrong TP plan. + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index cc5a564ab86f..730bf70db8ef 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -577,8 +577,10 @@ def forward( @auto_docstring class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py index 868018d744b5..41921eb080e6 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -70,22 +70,27 @@ class Glm4MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", # NOTE(3outeille): This needs to be right after down_proj in the dict. Otherwise, the pattern model.layers.*.mlp.experts will have priority over model.layers.*.mlp.experts.down_proj which will assign a wrong TP plan. + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py index a9518ed9c5d3..cc0c1fc67ffe 100644 --- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py @@ -53,21 +53,25 @@ class Glm4MoeLiteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py index 0b8ccc865775..4bab2f0316ae 100644 --- a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py @@ -651,8 +651,10 @@ def forward( @auto_docstring class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py index 1f65f44a525a..ea8e5f60e491 100644 --- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py @@ -61,21 +61,25 @@ class Glm4MoeLiteConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm4v/configuration_glm4v.py b/src/transformers/models/glm4v/configuration_glm4v.py index 4f151aa38156..dc993ab631d6 100644 --- a/src/transformers/models/glm4v/configuration_glm4v.py +++ b/src/transformers/models/glm4v/configuration_glm4v.py @@ -93,15 +93,22 @@ class Glm4vTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151552 diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py index a96fe6ace6b0..fdbf6473206a 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -134,15 +134,22 @@ class Glm4vTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151552 diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 0e4d6a9cb191..73971460b2d4 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -56,16 +56,23 @@ class Glm4vMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } @@ -98,6 +105,7 @@ class Glm4vMoeTextConfig(PreTrainedConfig): eos_token_id: int | list[int] | None = None pad_token_id: int | None = None base_config_key = "text_config" + ignore_keys_at_rope_validation = {"mrope_section"} router_aux_loss_coef: float = 0.0001 diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index b1c1293a00b9..289dac43e160 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -88,16 +88,23 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151424 diff --git a/src/transformers/models/glm_image/configuration_glm_image.py b/src/transformers/models/glm_image/configuration_glm_image.py index fcd302e35560..e022da23e861 100644 --- a/src/transformers/models/glm_image/configuration_glm_image.py +++ b/src/transformers/models/glm_image/configuration_glm_image.py @@ -106,15 +106,22 @@ class GlmImageTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 168064 diff --git a/src/transformers/models/glm_image/modular_glm_image.py b/src/transformers/models/glm_image/modular_glm_image.py index 273861a5b61e..8ce758601acd 100644 --- a/src/transformers/models/glm_image/modular_glm_image.py +++ b/src/transformers/models/glm_image/modular_glm_image.py @@ -133,6 +133,12 @@ class GlmImageTextConfig(Glm4vTextConfig): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 168064 max_position_embeddings: int = 131072 vision_vocab_size: int = 16512 diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 8f11f42794b3..31512bde976a 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -60,24 +60,27 @@ class GlmMoeDsaConfig(PreTrainedConfig): base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", "head_dim": "qk_rope_head_dim", diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 736dcdce32c3..1098096bd4e3 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -821,8 +821,10 @@ def forward( @auto_docstring class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index d909bb97e704..d45602619458 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -108,18 +108,21 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): base_model_tp_plan = { "layers.*.self_attn.q_b_proj": "colwise", - "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.kv_b_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } hidden_size: int = 6144 diff --git a/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/src/transformers/models/glm_ocr/configuration_glm_ocr.py index d08710c5e8b5..e1381b8bc79d 100644 --- a/src/transformers/models/glm_ocr/configuration_glm_ocr.py +++ b/src/transformers/models/glm_ocr/configuration_glm_ocr.py @@ -94,15 +94,22 @@ class GlmOcrTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 59392 diff --git a/src/transformers/models/glm_ocr/modular_glm_ocr.py b/src/transformers/models/glm_ocr/modular_glm_ocr.py index 782808bb69a9..2e9fc1c25a26 100644 --- a/src/transformers/models/glm_ocr/modular_glm_ocr.py +++ b/src/transformers/models/glm_ocr/modular_glm_ocr.py @@ -82,6 +82,12 @@ class GlmOcrTextConfig(Glm4vTextConfig): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 59392 hidden_size: int = 1024 intermediate_size: int = 4096 diff --git a/src/transformers/models/glmasr/modeling_glmasr.py b/src/transformers/models/glmasr/modeling_glmasr.py index e7173dc6dd8e..29b82adfc52c 100644 --- a/src/transformers/models/glmasr/modeling_glmasr.py +++ b/src/transformers/models/glmasr/modeling_glmasr.py @@ -373,6 +373,7 @@ class GlmAsrModelOutputWithPast(BaseModelOutputWithPast): ) class GlmAsrModel(GlmAsrPreTrainedModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None diff --git a/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/src/transformers/models/gpt_neox/configuration_gpt_neox.py index 256ddca26b49..abd0b4578a72 100644 --- a/src/transformers/models/gpt_neox/configuration_gpt_neox.py +++ b/src/transformers/models/gpt_neox/configuration_gpt_neox.py @@ -47,9 +47,18 @@ class GPTNeoXConfig(PreTrainedConfig): keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { "layers.*.attention.query_key_value": "colwise", - "layers.*.attention.dense": "rowwise", + "layers.*.attention.dense": "rowwise_allreduce", "layers.*.mlp.dense_h_to_4h": "colwise", - "layers.*.mlp.dense_4h_to_h": "rowwise", + "layers.*.mlp.dense_4h_to_h": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.dense_h_to_4h": "colwise", + "layers.*.mlp.dense_4h_to_h": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_in": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +67,12 @@ class GPTNeoXConfig(PreTrainedConfig): "final_layer_norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50432 hidden_size: int = 6144 num_hidden_layers: int = 44 diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py index 10e4b5922add..cfb81a96227d 100755 --- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -394,7 +394,7 @@ def set_input_embeddings(self, value): ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": "colwise_gather_output"} + _tp_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/modular_gpt_neox.py b/src/transformers/models/gpt_neox/modular_gpt_neox.py index f778501b7b38..cabda14021f0 100644 --- a/src/transformers/models/gpt_neox/modular_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modular_gpt_neox.py @@ -341,7 +341,7 @@ def forward( ) class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} - _tp_plan = {"embed_out": "colwise_gather_output"} + _tp_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index 47c029a5bca9..c7cf246fe259 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -32,13 +32,18 @@ class GptOssConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } base_model_ep_plan = { "layers.*.mlp.router": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } num_hidden_layers: int = 36 diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 484b071f01bc..ebb980bd8a9b 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -592,8 +592,10 @@ def load_balancing_loss_func( @auto_docstring class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/granite/configuration_granite.py b/src/transformers/models/granite/configuration_granite.py index e026cbbe5ff3..1464157b962d 100644 --- a/src/transformers/models/granite/configuration_granite.py +++ b/src/transformers/models/granite/configuration_granite.py @@ -50,10 +50,25 @@ class GraniteConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class GraniteConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 934345fe6723..5a266b3acc93 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -445,8 +445,10 @@ def forward( @auto_docstring class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/granite4_vision/configuration_granite4_vision.py b/src/transformers/models/granite4_vision/configuration_granite4_vision.py index 82c9e6765515..d3fb9129e5d2 100644 --- a/src/transformers/models/granite4_vision/configuration_granite4_vision.py +++ b/src/transformers/models/granite4_vision/configuration_granite4_vision.py @@ -53,10 +53,25 @@ class Granite4VisionTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +79,12 @@ class Granite4VisionTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 5fb53d6afe49..4f79272f164f 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -626,8 +626,10 @@ def load_balancing_loss_func( @auto_docstring class GraniteMoeForCausalLM(GraniteMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeConfig): super().__init__(config) diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 2e0926f3e5d4..929a82f5ed82 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1307,8 +1307,10 @@ def load_balancing_loss_func( @auto_docstring class GraniteMoeHybridForCausalLM(GraniteMoeHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeHybridConfig): super().__init__(config) diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 71f8c6eaff7d..8d0bc797a57b 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -695,8 +695,10 @@ def load_balancing_loss_func( @auto_docstring class GraniteMoeSharedForCausalLM(GraniteMoeSharedPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: GraniteMoeSharedConfig): super().__init__(config) diff --git a/src/transformers/models/helium/configuration_helium.py b/src/transformers/models/helium/configuration_helium.py index caa966b58a9d..5aca27c00069 100644 --- a/src/transformers/models/helium/configuration_helium.py +++ b/src/transformers/models/helium/configuration_helium.py @@ -44,10 +44,25 @@ class HeliumConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -55,6 +70,12 @@ class HeliumConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 48000 hidden_size: int = 2560 intermediate_size: int = 7040 diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index 8283fcb19e28..555fdc7592d3 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -423,8 +423,10 @@ def forward( @auto_docstring class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py index 97823eb79576..d1cd942ae268 100644 --- a/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py +++ b/src/transformers/models/higgs_audio_v2/configuration_higgs_audio_v2.py @@ -62,10 +62,25 @@ class HiggsAudioV2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +88,12 @@ class HiggsAudioV2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 3072 intermediate_size: int = 8192 diff --git a/src/transformers/models/hrm_text/configuration_hrm_text.py b/src/transformers/models/hrm_text/configuration_hrm_text.py index c8a328906b96..255391399960 100644 --- a/src/transformers/models/hrm_text/configuration_hrm_text.py +++ b/src/transformers/models/hrm_text/configuration_hrm_text.py @@ -65,12 +65,33 @@ class HrmTextConfig(PreTrainedConfig): **{f"{stack}.layers.*.mlp.up_proj": "colwise" for stack in ("L_module", "H_module")}, **{f"{stack}.layers.*.mlp.down_proj": "rowwise" for stack in ("L_module", "H_module")}, } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151808 hidden_size: int = 1536 intermediate_size: int = 4096 diff --git a/src/transformers/models/hrm_text/modeling_hrm_text.py b/src/transformers/models/hrm_text/modeling_hrm_text.py index 9e10bed4997e..731f3cba1ebb 100644 --- a/src/transformers/models/hrm_text/modeling_hrm_text.py +++ b/src/transformers/models/hrm_text/modeling_hrm_text.py @@ -555,8 +555,10 @@ def forward( @auto_docstring class HrmTextForCausalLM(HrmTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/hubert/modeling_hubert.py b/src/transformers/models/hubert/modeling_hubert.py index 44ee98aeffbf..6924271e4736 100755 --- a/src/transformers/models/hubert/modeling_hubert.py +++ b/src/transformers/models/hubert/modeling_hubert.py @@ -27,8 +27,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py index d1652d78cbbc..3b2386cf8ab3 100644 --- a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py +++ b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py @@ -461,8 +461,10 @@ def forward( @auto_docstring class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py index 19779da0528c..db33809260af 100644 --- a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py +++ b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py @@ -550,8 +550,10 @@ def forward( @auto_docstring class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py index e2eee94b118a..dddb7269f95e 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -54,18 +54,18 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +73,12 @@ class HYV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 120832 hidden_size: int = 4096 intermediate_size: int = 13312 diff --git a/src/transformers/models/hy_v3/modeling_hy_v3.py b/src/transformers/models/hy_v3/modeling_hy_v3.py index f2d64736b4f0..6ade5c2d278d 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -544,8 +544,10 @@ def forward( @auto_docstring class HYV3ForCausalLM(HYV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: HYV3Config): super().__init__(config) diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py index fa0931435197..5c9547a2fe0e 100644 --- a/src/transformers/models/hy_v3/modular_hy_v3.py +++ b/src/transformers/models/hy_v3/modular_hy_v3.py @@ -79,18 +79,18 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -98,6 +98,12 @@ class HYV3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 120832 hidden_size: int = 4096 intermediate_size: int = 13312 diff --git a/src/transformers/models/hyperclovax/configuration_hyperclovax.py b/src/transformers/models/hyperclovax/configuration_hyperclovax.py index 430a56bf0249..1090a4e00fda 100644 --- a/src/transformers/models/hyperclovax/configuration_hyperclovax.py +++ b/src/transformers/models/hyperclovax/configuration_hyperclovax.py @@ -65,10 +65,25 @@ class HyperCLOVAXConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +91,12 @@ class HyperCLOVAXConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/hyperclovax/modeling_hyperclovax.py b/src/transformers/models/hyperclovax/modeling_hyperclovax.py index 3608d215bfa9..a69fb46b121a 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -452,8 +452,10 @@ def forward( @auto_docstring class HyperCLOVAXForCausalLM(HyperCLOVAXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/jais2/configuration_jais2.py b/src/transformers/models/jais2/configuration_jais2.py index 7139661d7575..6fa9783effd3 100644 --- a/src/transformers/models/jais2/configuration_jais2.py +++ b/src/transformers/models/jais2/configuration_jais2.py @@ -50,9 +50,23 @@ class Jais2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,6 +74,12 @@ class Jais2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 150272 hidden_size: int = 3328 intermediate_size: int = 26624 diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 5e6a37c0172d..1fb6c19d9044 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -397,8 +397,10 @@ def forward( @auto_docstring class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/jais2/modular_jais2.py b/src/transformers/models/jais2/modular_jais2.py index 3fbdf2c8cd46..c760e70c4b15 100644 --- a/src/transformers/models/jais2/modular_jais2.py +++ b/src/transformers/models/jais2/modular_jais2.py @@ -34,9 +34,23 @@ class Jais2Config(LlamaConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } vocab_size: int = 150272 diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index d9e3ff7b84ef..ba5a759c4636 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -843,8 +843,10 @@ def load_balancing_loss_func( @auto_docstring class JambaForCausalLM(JambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config: JambaConfig): super().__init__(config) diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py index b16274332baf..fbe51b44259d 100644 --- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py @@ -874,8 +874,10 @@ def forward( @auto_docstring class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} _keep_in_fp32_modules_strict = ["codec_model"] output_modalities = ("audio", "text") diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py index 33f939f6db43..50e011ef08be 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -61,18 +61,18 @@ class LagunaConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -80,6 +80,12 @@ class LagunaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 2048 intermediate_size: int = 8192 diff --git a/src/transformers/models/laguna/modeling_laguna.py b/src/transformers/models/laguna/modeling_laguna.py index aa4060e77f5f..410fd9a04258 100644 --- a/src/transformers/models/laguna/modeling_laguna.py +++ b/src/transformers/models/laguna/modeling_laguna.py @@ -672,8 +672,10 @@ def load_balancing_loss_func( @auto_docstring class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/laguna/modular_laguna.py b/src/transformers/models/laguna/modular_laguna.py index 945cd40a99b2..7587827acb2c 100644 --- a/src/transformers/models/laguna/modular_laguna.py +++ b/src/transformers/models/laguna/modular_laguna.py @@ -79,18 +79,18 @@ class LagunaConfig(Qwen2MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", } vocab_size: int = 100352 diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index ef753e3b2893..95af77e0141b 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -540,8 +540,10 @@ def forward( @auto_docstring class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 0369ae31b8ae..dc97dfb11454 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -630,8 +630,10 @@ def forward( @auto_docstring class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py index 6960a6970592..0f3a0e75f7f6 100644 --- a/src/transformers/models/llama/configuration_llama.py +++ b/src/transformers/models/llama/configuration_llama.py @@ -50,10 +50,25 @@ class LlamaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class LlamaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 9d659c7c6f08..655e542358f5 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -428,8 +428,10 @@ def forward( @auto_docstring class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 79cfd063f4d4..781de8a57fd0 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -45,10 +45,10 @@ class Llama4VisionConfig(PreTrainedConfig): "model.layers.*.self_attn.q_proj": "colwise", "model.layers.*.self_attn.k_proj": "colwise", "model.layers.*.self_attn.v_proj": "colwise", - "model.layers.*.self_attn.o_proj": "rowwise", + "model.layers.*.self_attn.o_proj": "rowwise_allreduce", "vision_adapter.mlp.fc1": "colwise", - "vision_adapter.mlp.fc2": "rowwise", - "patch_embedding.linear": "colwise_gather_output", + "vision_adapter.mlp.fc2": "rowwise_allreduce", + "patch_embedding.linear": "colwise_allgather", } model_type = "llama4_vision_model" base_config_key = "vision_config" @@ -113,26 +113,24 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.feed_forward.shared_expert.gate_proj": "colwise", "layers.*.feed_forward.shared_expert.up_proj": "colwise", - "layers.*.feed_forward.shared_expert.down_proj": "rowwise", - "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise", # row because not linear - "layers.*.feed_forward.experts.down_proj": "colwise", # col because not linear + "layers.*.feed_forward.shared_expert.down_proj": "rowwise_allreduce", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", - "layers.*.feed_forward.down_proj": "rowwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", } base_model_ep_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", - "layers.*.feed_forward.down_proj": "rowwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", "layers.*.feed_forward.router": "ep_router", } @@ -179,7 +177,7 @@ def __post_init__(self, **kwargs): default_no_rope_layers = [ int((layer_idx + 1) % self.no_rope_layer_interval != 0) for layer_idx in range(self.num_hidden_layers) ] - self.no_rope_layers = self.no_rope_layers if self.no_rope_layers else default_no_rope_layers + self.no_rope_layers = self.no_rope_layers or default_no_rope_layers self.head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads self.moe_layers = ( diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index e064e399e6c0..13dd361ecb49 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -590,7 +590,9 @@ class Llama4ForCausalLM(Llama4PreTrainedModel, GenerationMixin): _no_split_modules = ["Llama4TextDecoderLayer"] base_model_prefix = "language_model" _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} config: Llama4TextConfig def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index 39e5a03338d8..ffeb6de9d01f 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -55,16 +55,11 @@ class LongcatFlashConfig(PreTrainedConfig): default_theta = 10000000.0 base_model_tp_plan = { "layers.*.self_attn.*.q_b_proj": "colwise", - "layers.*.self_attn.*.kv_a_proj_with_mqa": "mla_kv_a_proj", "layers.*.self_attn.*.kv_b_proj": "colwise", - "layers.*.self_attn.*.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts.identity_expert": "moe_identity_expert", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.*.o_proj": "rowwise_allreduce", "layers.*.mlps.*.gate_proj": "colwise", "layers.*.mlps.*.up_proj": "colwise", - "layers.*.mlps.*.down_proj": "rowwise", + "layers.*.mlps.*.down_proj": "rowwise_allreduce", } base_model_pp_plan = { @@ -73,6 +68,12 @@ class LongcatFlashConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 131072 hidden_size: int = 6144 num_hidden_layers: int = 56 diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index d5ac6e237742..4e7b3ad287cd 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -650,8 +650,10 @@ def forward( @auto_docstring class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] def __init__(self, config): diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 9a3a0023725e..d60e6d0c2849 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -65,16 +65,21 @@ class MiniMaxConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = {"num_experts": "num_local_experts"} vocab_size: int = 32000 diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 69497f83cad8..7da5e5867ce2 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -789,8 +789,10 @@ def load_balancing_loss_func( @auto_docstring class MiniMaxForCausalLM(MiniMaxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 0bd400458129..8442eebd53e7 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -91,16 +91,21 @@ class MiniMaxConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = {"num_experts": "num_local_experts"} vocab_size: int = 32000 diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 75acb8b755d7..6193523d95b3 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -48,19 +48,37 @@ class MiniMaxM2Config(PreTrainedConfig): model_type = "minimax_m2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", - "layers.*.self_attn.k_proj": "colwise_gather_output", - "layers.*.self_attn.v_proj": "colwise_gather_output", - "layers.*.self_attn.o_proj": "rowwise_split_input", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_experts": "num_local_experts", } diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index d19274262810..370ea88c971e 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -588,8 +588,10 @@ def load_balancing_loss_func( @auto_docstring class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index a9938a555c62..ca1b29125ccf 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -67,19 +67,37 @@ class MiniMaxM2Config(PreTrainedConfig): model_type = "minimax_m2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", - "layers.*.self_attn.k_proj": "colwise_gather_output", - "layers.*.self_attn.v_proj": "colwise_gather_output", - "layers.*.self_attn.o_proj": "rowwise_split_input", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_experts": "num_local_experts", } diff --git a/src/transformers/models/ministral/configuration_ministral.py b/src/transformers/models/ministral/configuration_ministral.py index 4ff445e6808a..fa0de0f0671e 100644 --- a/src/transformers/models/ministral/configuration_ministral.py +++ b/src/transformers/models/ministral/configuration_ministral.py @@ -52,10 +52,25 @@ class MinistralConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -63,6 +78,12 @@ class MinistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index af4f7fbeae59..d2faea85b8cc 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -429,8 +429,10 @@ def forward( @auto_docstring class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/ministral3/configuration_ministral3.py b/src/transformers/models/ministral3/configuration_ministral3.py index 76d55710d5ae..d486d2efe796 100644 --- a/src/transformers/models/ministral3/configuration_ministral3.py +++ b/src/transformers/models/ministral3/configuration_ministral3.py @@ -58,16 +58,37 @@ class Ministral3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } ignore_keys_at_rope_validation = {"llama_4_scaling_beta", "max_position_embeddings"} vocab_size: int = 131072 diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index 6aacf4c8ce3a..066a71997913 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -412,8 +412,10 @@ def forward( @auto_docstring class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mistral/configuration_mistral.py b/src/transformers/models/mistral/configuration_mistral.py index c57193f58d7b..0451d2f958a3 100644 --- a/src/transformers/models/mistral/configuration_mistral.py +++ b/src/transformers/models/mistral/configuration_mistral.py @@ -49,10 +49,25 @@ class MistralConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,6 +75,12 @@ class MistralConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index b79dea36c9e9..d94ff3d6a312 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -401,8 +401,10 @@ def forward( @auto_docstring class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index 0e16e0a14f45..c8218ffa5c57 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -47,21 +47,25 @@ class Mistral4Config(PreTrainedConfig): model_type = "mistral4" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = { "num_local_experts": "n_routed_experts", } diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 006ddad187bf..4473107418ee 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -640,8 +640,10 @@ def forward( @auto_docstring class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 240f24411031..41beb623719c 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -42,20 +42,41 @@ class MixtralConfig(PreTrainedConfig): model_type = "mixtral" keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.0 + # TP plan (for inference/generation). base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + + # TP + Sequence Parallelism plan (for training). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + # FSDP2 plan (see Qwen3Config.base_model_fsdp_plan for shape rationale). + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = {"num_experts": "num_local_experts"} vocab_size: int = 32000 diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 991851dbadd3..df71b44ebac2 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -580,8 +580,10 @@ def load_balancing_loss_func( @auto_docstring class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py index 139e580fbca7..438005523f7e 100644 --- a/src/transformers/models/mixtral/modular_mixtral.py +++ b/src/transformers/models/mixtral/modular_mixtral.py @@ -335,6 +335,7 @@ def forward( class MixtralForCausalLM(MistralForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/musicflamingo/modeling_musicflamingo.py b/src/transformers/models/musicflamingo/modeling_musicflamingo.py index cf0cb468ec1f..f59f3422887d 100644 --- a/src/transformers/models/musicflamingo/modeling_musicflamingo.py +++ b/src/transformers/models/musicflamingo/modeling_musicflamingo.py @@ -213,6 +213,7 @@ def apply_rotary_time_emb(hidden_states, cos, sin): ) class MusicFlamingoModel(MusicFlamingoPreTrainedModel): _tp_plan = None + _sp_plan = None _pp_plan = None _keep_in_fp32_modules_strict = None diff --git a/src/transformers/models/nanochat/configuration_nanochat.py b/src/transformers/models/nanochat/configuration_nanochat.py index 24a0ab7b6d09..12ee29c11ef6 100644 --- a/src/transformers/models/nanochat/configuration_nanochat.py +++ b/src/transformers/models/nanochat/configuration_nanochat.py @@ -46,10 +46,11 @@ class NanoChatConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.fc1": "colwise", - "layers.*.mlp.fc2": "rowwise", + "layers.*.mlp.fc2": "rowwise_allreduce", } + base_model_sp_plan = None vocab_size: int = 50304 hidden_size: int = 768 diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index 9205b89cd360..a191eb2d5425 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -432,8 +432,10 @@ def forward( @auto_docstring class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index 713cc29b81eb..460cf321150f 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -198,7 +198,9 @@ def forward( @auto_docstring class NanoChatForCausalLM(Gemma2ForCausalLM): - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: r""" diff --git a/src/transformers/models/nllb_moe/modeling_nllb_moe.py b/src/transformers/models/nllb_moe/modeling_nllb_moe.py index d4ef4927ce67..7370722e38b0 100644 --- a/src/transformers/models/nllb_moe/modeling_nllb_moe.py +++ b/src/transformers/models/nllb_moe/modeling_nllb_moe.py @@ -22,9 +22,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/olmo/configuration_olmo.py b/src/transformers/models/olmo/configuration_olmo.py index 186cc3a704fb..f9f32c1ed15d 100644 --- a/src/transformers/models/olmo/configuration_olmo.py +++ b/src/transformers/models/olmo/configuration_olmo.py @@ -53,10 +53,25 @@ class OlmoConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +79,12 @@ class OlmoConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index 0a886321bfc3..dad24db7ce06 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -425,8 +425,10 @@ def forward( @auto_docstring class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo2/configuration_olmo2.py b/src/transformers/models/olmo2/configuration_olmo2.py index f879c0b8367f..719db72c42e1 100644 --- a/src/transformers/models/olmo2/configuration_olmo2.py +++ b/src/transformers/models/olmo2/configuration_olmo2.py @@ -53,13 +53,30 @@ class Olmo2Config(PreTrainedConfig): model_type = "olmo2" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -67,6 +84,12 @@ class Olmo2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index 9901212ff737..12d31e551870 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -429,8 +429,10 @@ def forward( @auto_docstring class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo2/modular_olmo2.py b/src/transformers/models/olmo2/modular_olmo2.py index 4ac66c2e4608..3407011afffe 100644 --- a/src/transformers/models/olmo2/modular_olmo2.py +++ b/src/transformers/models/olmo2/modular_olmo2.py @@ -67,13 +67,30 @@ class Olmo2Config(OlmoConfig): model_type = "olmo2" base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo3/configuration_olmo3.py b/src/transformers/models/olmo3/configuration_olmo3.py index 2f45be450a0b..cf81a4b3e332 100644 --- a/src/transformers/models/olmo3/configuration_olmo3.py +++ b/src/transformers/models/olmo3/configuration_olmo3.py @@ -48,13 +48,30 @@ class Olmo3Config(PreTrainedConfig): model_type = "olmo3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -62,6 +79,12 @@ class Olmo3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 50304 hidden_size: int = 4096 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index 78ade3570f97..8325b651312a 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -433,8 +433,10 @@ def forward( @auto_docstring class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo3/modular_olmo3.py b/src/transformers/models/olmo3/modular_olmo3.py index f5934880c5f5..b587e4a2de30 100644 --- a/src/transformers/models/olmo3/modular_olmo3.py +++ b/src/transformers/models/olmo3/modular_olmo3.py @@ -63,13 +63,30 @@ class Olmo3Config(Olmo2Config): model_type = "olmo3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py index 0f7c32e8799e..106fb2010bfc 100644 --- a/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/configuration_olmo_hybrid.py @@ -73,20 +73,27 @@ class OlmoHybridConfig(PreTrainedConfig): model_type = "olmo_hybrid" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 100352 hidden_size: int = 3840 intermediate_size: int = 11008 diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index 5f557dccc12d..081149371067 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1045,8 +1045,10 @@ def _update_linear_attn_mask(self, attention_mask, past_key_values): @auto_docstring class OlmoHybridForCausalLM(OlmoHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py index ef805ce9498a..df028e534aa3 100644 --- a/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py @@ -120,14 +120,15 @@ class OlmoHybridConfig(LlamaConfig): model_type = "olmo_hybrid" base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.k_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.v_proj": "colwise_gather_output", # we need to replicate here due to the added norm on q and k - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the added norm on q and k + "layers.*.self_attn.q_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None vocab_size: int = 100352 hidden_size: int = 3840 diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 16bedbe698f8..feedb02bfda4 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -46,13 +46,30 @@ class OlmoeConfig(PreTrainedConfig): # Default tensor parallel plan for base model `Olmoe` base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", # due to the norm, we have to gather - "layers.*.self_attn.k_proj": "colwise_gather_output", # due to the norm, we have to gather - "layers.*.self_attn.v_proj": "colwise_gather_output", # due to the norm, we have to gather - "layers.*.self_attn.o_proj": "rowwise_split_input", # due to the norm, we have to gather - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.k_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.v_proj": "colwise_allgather", # due to the norm, we have to gather + "layers.*.self_attn.o_proj": "vocab_allreduce", # due to the norm, we have to gather + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } vocab_size: int = 50304 diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 5d89ec741529..10950656c462 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -604,8 +604,10 @@ def load_balancing_loss_func( @auto_docstring class OlmoeForCausalLM(OlmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py index e7aaefde4bca..5f2933c91984 100644 --- a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py @@ -56,16 +56,23 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } base_model_ep_plan = { "layers.*.mlp.router": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } num_hidden_layers: int = 8 num_local_experts: int = 128 + vocab_size: int = 200064 hidden_size: int = 640 intermediate_size: int = 640 diff --git a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py index 422235d9da91..7a438fedd042 100644 --- a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py @@ -72,6 +72,13 @@ @strict class OpenAIPrivacyFilterConfig(GptOssConfig): model_type = "openai_privacy_filter" + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 200064 hidden_size: int = 640 intermediate_size: int = 640 diff --git a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py index 343a22ade814..b9bef13987ac 100644 --- a/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py @@ -98,10 +98,25 @@ class PaddleOCRTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -109,6 +124,12 @@ class PaddleOCRTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 103424 hidden_size: int = 1024 intermediate_size: int = 3072 diff --git a/src/transformers/models/phi/configuration_phi.py b/src/transformers/models/phi/configuration_phi.py index 2a65f0f16aec..bee22ad6a0b7 100644 --- a/src/transformers/models/phi/configuration_phi.py +++ b/src/transformers/models/phi/configuration_phi.py @@ -49,9 +49,23 @@ class PhiConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.dense": "rowwise", + "layers.*.self_attn.dense": "rowwise_allreduce", "layers.*.mlp.fc1": "colwise", - "layers.*.mlp.fc2": "rowwise", + "layers.*.mlp.fc2": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.dense": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.fc1": "colwise", + "layers.*.mlp.fc2": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -60,14 +74,20 @@ class PhiConfig(PreTrainedConfig): "final_layernorm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 51200 hidden_size: int = 2048 intermediate_size: int = 8192 num_hidden_layers: int = 24 num_attention_heads: int = 32 num_key_value_heads: int | None = None - resid_pdrop: float | int = 0.0 - embd_pdrop: float | int = 0.0 + resid_pdrop: float = 0.0 + embd_pdrop: float = 0.0 attention_dropout: float | int | None = 0.0 hidden_act: str = "gelu_new" max_position_embeddings: int = 2048 diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index e3f97a01ee4c..0acf6a8060aa 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -406,8 +406,10 @@ def forward( @auto_docstring class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phi3/configuration_phi3.py b/src/transformers/models/phi3/configuration_phi3.py index f85502f8205d..d1288d411863 100644 --- a/src/transformers/models/phi3/configuration_phi3.py +++ b/src/transformers/models/phi3/configuration_phi3.py @@ -47,10 +47,22 @@ class Phi3Config(PreTrainedConfig): model_type = "phi3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.qkv_proj": "colwise_gather_output", # we need to replicate here due to the slicing of qkv - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the slicing of qkv - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.qkv_proj": "colwise_allgather", # we need to replicate here due to the slicing of qkv + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the slicing of qkv + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.qkv_proj": "colwise_allgather", # fused qkv needs full tensor for slicing + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": "vocab_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +70,12 @@ class Phi3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32064 hidden_size: int = 3072 intermediate_size: int = 8192 diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index b07735f8a2e6..f0d6726b73d9 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -432,8 +432,10 @@ def forward( @auto_docstring class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py index b73741305566..264119a82bec 100644 --- a/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/configuration_phi4_multimodal.py @@ -181,10 +181,22 @@ class Phi4MultimodalConfig(PreTrainedConfig): model_type = "phi4_multimodal" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.qkv_proj": "colwise_gather_output", # we need to replicate here due to the slicing of qkv - "layers.*.self_attn.o_proj": "rowwise_split_input", # input is replicated due to the slicing of qkv - "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation - "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation + "layers.*.self_attn.qkv_proj": "colwise_allgather", # we need to replicate here due to the slicing of qkv + "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the slicing of qkv + "layers.*.mlp.gate_up_proj": "colwise_allgather", # we need to replicate here due to the `chunk` operation + "layers.*.mlp.down_proj": "vocab_allreduce", # input is replicated due to the `chunk` operation + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.qkv_proj": "colwise_allgather", # fused qkv needs full tensor for slicing + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_up_proj": "colwise_allgather", # fused gate/up needs full tensor for chunk + "layers.*.mlp.down_proj": "vocab_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -192,6 +204,12 @@ class Phi4MultimodalConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 200064 hidden_size: int = 3072 intermediate_size: int = 8192 diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 9e6c0339098d..383445a29f8b 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1596,8 +1596,10 @@ def forward( @auto_docstring class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/phimoe/configuration_phimoe.py b/src/transformers/models/phimoe/configuration_phimoe.py index e20f94085be0..0545ffe5ea70 100644 --- a/src/transformers/models/phimoe/configuration_phimoe.py +++ b/src/transformers/models/phimoe/configuration_phimoe.py @@ -47,6 +47,11 @@ class PhimoeConfig(PreTrainedConfig): model_type = "phimoe" keys_to_ignore_at_inference = ["past_key_values"] default_theta = 1000000.0 + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } vocab_size: int = 32064 hidden_size: int = 4096 diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 23bc944c522a..d933a6ec69d9 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -772,8 +772,10 @@ def load_balancing_loss_func( @auto_docstring class PhimoeForCausalLM(PhimoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/pi0/modeling_pi0.py b/src/transformers/models/pi0/modeling_pi0.py index bd86b6f10644..da1e4a9c8b35 100644 --- a/src/transformers/models/pi0/modeling_pi0.py +++ b/src/transformers/models/pi0/modeling_pi0.py @@ -222,7 +222,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": "colwise_gather_output"} + _tp_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/pi0/modular_pi0.py b/src/transformers/models/pi0/modular_pi0.py index 337ae9406bb4..53472a40ccfe 100644 --- a/src/transformers/models/pi0/modular_pi0.py +++ b/src/transformers/models/pi0/modular_pi0.py @@ -482,7 +482,7 @@ def forward( class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" - _tp_plan = {"action_out_proj": "colwise_gather_output"} + _tp_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/qwen2/configuration_qwen2.py b/src/transformers/models/qwen2/configuration_qwen2.py index ae41af7b211f..9106216fb13d 100644 --- a/src/transformers/models/qwen2/configuration_qwen2.py +++ b/src/transformers/models/qwen2/configuration_qwen2.py @@ -47,10 +47,25 @@ class Qwen2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +73,12 @@ class Qwen2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 4096 intermediate_size: int = 22016 diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 9263e1d42937..7dcc274a2cf2 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -416,8 +416,10 @@ def forward( @auto_docstring class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py index 081823bf222f..fc90da09bbbf 100644 --- a/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/configuration_qwen2_5_omni.py @@ -157,16 +157,23 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py index acebc79a3c94..e9eeefc9b9f8 100644 --- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py @@ -259,16 +259,23 @@ class Qwen2_5OmniTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py index 385f2ecad057..da05af794ed5 100644 --- a/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py @@ -91,16 +91,22 @@ class Qwen2_5_VLTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index 7f961976b58c..329a6ef2a058 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -57,10 +57,10 @@ class Qwen2MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -68,6 +68,12 @@ class Qwen2MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 5632 diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index d4150d0a74d7..f7142b4650fb 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -617,8 +617,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen2MoeForCausalLM(Qwen2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py index deb615c9e7b6..9c5b65f47896 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -230,7 +230,9 @@ def forward( class Qwen2MoeForCausalLM(MixtralForCausalLM, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py index 65680d492a78..bb6871f46a05 100644 --- a/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py @@ -68,16 +68,22 @@ class Qwen2VLTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 152064 diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index 07ad0bb24b33..ff0cd8f63a23 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -41,17 +41,41 @@ class Qwen3Config(PreTrainedConfig): model_type = "qwen3" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `Qwen3` + # TP plan (for inference/generation). + # All activations are plain tensors — compatible with KV cache and autoregressive + # decode (seq_len=1). Each rank holds a full copy of activations between layers. base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + # TP + Sequence Parallelism plan (for training). + # Activations between layers are sharded on the sequence dimension (Shard(1)), + # reducing per-rank activation memory by tp_size. In exchange, extra collectives + # (all-gather before attention/MLP, reduce-scatter after) are needed. + # Not compatible with autoregressive decode (because seq_len=1 can't be split across ranks) + # or KV cache (which stores plain tensors). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -59,6 +83,15 @@ class Qwen3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) + # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # into a single fully_shard([...]) call at apply time so they share one all-gather. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 4096 intermediate_size: int = 22016 diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 91715a33cf9d..2ba669225879 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -441,8 +441,10 @@ def forward( @auto_docstring class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3/modular_qwen3.py b/src/transformers/models/qwen3/modular_qwen3.py index 73cde6d89a7a..ea6ddb2f7499 100644 --- a/src/transformers/models/qwen3/modular_qwen3.py +++ b/src/transformers/models/qwen3/modular_qwen3.py @@ -108,6 +108,8 @@ def forward( class Qwen3ForCausalLM(Qwen2ForCausalLM): + _fsdp_plan = {"lm_head": "keep_full_weight"} + def forward( self, **super_kwargs: Unpack[TransformersKwargs], diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index ae9eb8f86c6d..04877a3bbdae 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -60,12 +60,10 @@ class Qwen3_5TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -73,6 +71,12 @@ class Qwen3_5TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 248320 hidden_size: int = 4096 intermediate_size: int = 12288 diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 587c7cf362fb..b97168db0a30 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1629,8 +1629,10 @@ def forward( @auto_docstring class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 91457c58f3ce..82d045c0a28b 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -96,12 +96,10 @@ class Qwen3_5TextConfig(Qwen3NextConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} diff --git a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py index f6f9594e0d73..72c2dacfd2c0 100644 --- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -60,15 +60,11 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +72,12 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 248320 hidden_size: int = 2048 num_hidden_layers: int = 40 diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index dbf459287d0e..295341bdd7b3 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -1831,8 +1831,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5MoeTextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] @@ -1938,7 +1940,9 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False config: Qwen3_5MoeConfig - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py index 70469ea91bc7..58591c6f2648 100644 --- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py @@ -89,15 +89,11 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", } ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} @@ -248,7 +244,9 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _fsdp_plan = {"lm_head": "keep_full_weight"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs): r""" diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index 9a7d4b4c8b5b..4395b7c622bd 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -57,15 +57,29 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } # Expert-only EP plan: only shards MoE experts, not attention. # Attention is left unsharded — FSDP2 handles attention weight distribution. @@ -74,7 +88,7 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -82,6 +96,12 @@ class Qwen3MoeConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 6144 diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index ddf84fc575b7..012e0cac7243 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -609,8 +609,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index bf26179ff3fd..b499291d70e2 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -62,18 +62,14 @@ class Qwen3NextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -81,6 +77,12 @@ class Qwen3NextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 intermediate_size: int = 5632 diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index df167395975b..faca94ba5c43 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1088,8 +1088,10 @@ def load_balancing_loss_func( @auto_docstring class Qwen3NextForCausalLM(Qwen3NextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py index 4e9bdd35f21f..3c65badeafb4 100644 --- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py @@ -138,18 +138,23 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section", "interleaved", "mrope_interleaved"} vocab_size: int = 3584 @@ -263,17 +268,41 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): model_type = "qwen3_omni_moe_talker_code_predictor" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `Qwen3OmniMoeTalkerCodePredictor` + # TP plan (for inference/generation). + # All activations are plain tensors — compatible with KV cache and autoregressive + # decode (seq_len=1). Each rank holds a full copy of activations between layers. base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + + # TP + Sequence Parallelism plan (for training). + # Activations between layers are sharded on the sequence dimension (Shard(1)), + # reducing per-rank activation memory by tp_size. In exchange, extra collectives + # (all-gather before attention/MLP, reduce-scatter after) are needed. + # Not compatible with autoregressive decode (because seq_len=1 can't be split across ranks) + # or KV cache (which stores plain tensors). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -281,6 +310,15 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) + # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # into a single fully_shard([...]) call at apply time so they share one all-gather. + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 2048 hidden_size: int = 1024 intermediate_size: int = 3072 @@ -357,21 +395,35 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -379,6 +431,12 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 3072 hidden_size: int = 1024 intermediate_size: int = 2048 diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py index 05f310a75b11..83875dc88b5d 100644 --- a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py @@ -2637,8 +2637,10 @@ def get_input_embeddings(self): @auto_docstring class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3OmniMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerCodePredictorConfig base_model_prefix = "talker.code_predictor" _can_record_outputs = { @@ -3021,8 +3023,10 @@ def get_input_embeddings(self): @auto_docstring class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": "colwise_gather_output"} + _tp_plan = {"codec_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" _no_split_modules = ["Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration"] diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index 9f8351d3f8c4..dd98da222cf3 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -280,18 +280,23 @@ class Qwen3OmniMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + ignore_keys_at_rope_validation = {"mrope_section", "interleaved", "mrope_interleaved"} vocab_size: int = 3584 @@ -394,7 +399,7 @@ class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } vocab_size: int = 3072 @@ -1659,7 +1664,7 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} - _tp_plan = {"codec_head": "colwise_gather_output"} + _tp_plan = {"codec_head": "colwise_allgather"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index 6825f5dda929..f49d389a593f 100644 --- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py @@ -59,16 +59,34 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -76,6 +94,12 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 151936 hidden_size: int = 2048 diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py index 11534b395773..555e3d63ccb3 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -86,16 +86,16 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py index e3f2d7a6ed26..10b31ee9c3c9 100755 --- a/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py +++ b/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py @@ -24,9 +24,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py index a2b292555fc8..48e85b6499c5 100644 --- a/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py +++ b/src/transformers/models/seamless_m4t_v2/modeling_seamless_m4t_v2.py @@ -24,9 +24,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/seed_oss/configuration_seed_oss.py b/src/transformers/models/seed_oss/configuration_seed_oss.py index b1221fcf53ce..363b393dbc4a 100644 --- a/src/transformers/models/seed_oss/configuration_seed_oss.py +++ b/src/transformers/models/seed_oss/configuration_seed_oss.py @@ -48,10 +48,25 @@ class SeedOssConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -59,6 +74,12 @@ class SeedOssConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 155136 hidden_size: int = 4096 intermediate_size: int = 27648 diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 1ebc8f10a272..841765a7c928 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -429,8 +429,10 @@ def forward( @auto_docstring class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/sew/modeling_sew.py b/src/transformers/models/sew/modeling_sew.py index ea499e63289a..30e527b4c985 100644 --- a/src/transformers/models/sew/modeling_sew.py +++ b/src/transformers/models/sew/modeling_sew.py @@ -28,8 +28,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, CausalLMOutput, SequenceClassifierOutput diff --git a/src/transformers/models/sew/modular_sew.py b/src/transformers/models/sew/modular_sew.py index 312419793a34..f3db3c1bcaaa 100644 --- a/src/transformers/models/sew/modular_sew.py +++ b/src/transformers/models/sew/modular_sew.py @@ -20,8 +20,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_outputs import BaseModelOutput from ...modeling_utils import PreTrainedModel from ...utils import auto_docstring diff --git a/src/transformers/models/smollm3/configuration_smollm3.py b/src/transformers/models/smollm3/configuration_smollm3.py index f48c979a1dd3..9a14e834db9e 100644 --- a/src/transformers/models/smollm3/configuration_smollm3.py +++ b/src/transformers/models/smollm3/configuration_smollm3.py @@ -58,10 +58,25 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -69,6 +84,12 @@ class SmolLM3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 2048 intermediate_size: int = 11008 diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index 8d911e414b0f..d82882382bb9 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -445,8 +445,10 @@ def forward( @auto_docstring class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/smollm3/modular_smollm3.py b/src/transformers/models/smollm3/modular_smollm3.py index f75017ad2645..69c9d4b51444 100644 --- a/src/transformers/models/smollm3/modular_smollm3.py +++ b/src/transformers/models/smollm3/modular_smollm3.py @@ -74,10 +74,25 @@ class SmolLM3Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -85,6 +100,12 @@ class SmolLM3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 128256 hidden_size: int = 2048 intermediate_size: int = 11008 diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index ac0016aa7791..5d9aec3ac7a0 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -41,16 +41,21 @@ class SolarOpenConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + attribute_map = { "num_local_experts": "n_routed_experts", } @@ -81,6 +86,19 @@ class SolarOpenConfig(PreTrainedConfig): eos_token_id: int | list[int] | None = None pad_token_id: int | None = None default_theta = 1_000_000.0 + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } head_dim: int = 128 def __post_init__(self, **kwargs): diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index 0eb50021ecd6..53abac7c7d07 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -552,8 +552,10 @@ def forward( @auto_docstring class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py index 90d4f0c389c0..48bba9439e08 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -47,10 +47,27 @@ class SolarOpenConfig(Glm4MoeConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", } vocab_size: int = 196608 diff --git a/src/transformers/models/speecht5/modeling_speecht5.py b/src/transformers/models/speecht5/modeling_speecht5.py index f75fa9dcdcd9..e5d7d706a9ab 100644 --- a/src/transformers/models/speecht5/modeling_speecht5.py +++ b/src/transformers/models/speecht5/modeling_speecht5.py @@ -23,9 +23,9 @@ from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache +from ...distributed.fsdp import is_fsdp_managed_module from ...generation import GenerationMixin from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/starcoder2/configuration_starcoder2.py b/src/transformers/models/starcoder2/configuration_starcoder2.py index 59efa94fc5f4..88137e8f362f 100644 --- a/src/transformers/models/starcoder2/configuration_starcoder2.py +++ b/src/transformers/models/starcoder2/configuration_starcoder2.py @@ -48,9 +48,23 @@ class Starcoder2Config(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.c_fc": "colwise", - "layers.*.mlp.c_proj": "rowwise", + "layers.*.mlp.c_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.c_fc": "colwise", + "layers.*.mlp.c_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -58,6 +72,12 @@ class Starcoder2Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 49152 hidden_size: int = 3072 intermediate_size: int = 12288 diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index 8b89a1d1745c..baabf385da59 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -409,8 +409,10 @@ def forward( @auto_docstring class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py index 9de40c832259..8c313bae875e 100644 --- a/src/transformers/models/t5gemma/configuration_t5gemma.py +++ b/src/transformers/models/t5gemma/configuration_t5gemma.py @@ -54,10 +54,25 @@ class T5GemmaModuleConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -65,6 +80,12 @@ class T5GemmaModuleConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index 074877e17f21..ceafdfa43d98 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -946,7 +946,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py index 553cfdb41c2a..f6a6595f9dbc 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -85,6 +85,12 @@ class T5GemmaModuleConfig(Gemma2Config): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + is_decoder: bool = False use_bidirectional_attention = AttributeError() @@ -785,7 +791,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index d9a9a3f5769f..b3c2ada26f35 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -48,12 +48,27 @@ class T5Gemma2TextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -61,6 +76,12 @@ class T5Gemma2TextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_208 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -85,6 +106,7 @@ class T5Gemma2TextConfig(PreTrainedConfig): layer_types: list[str] | None = None final_logit_softcapping: float | None = None attn_logit_softcapping: float | None = None + default_theta = {"global": 1_000_000.0, "local": 10_000.0} def __post_init__(self, **kwargs): @@ -223,12 +245,27 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -236,6 +273,12 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 262_208 hidden_size: int = 2304 intermediate_size: int = 9216 @@ -260,6 +303,7 @@ class T5Gemma2DecoderConfig(PreTrainedConfig): layer_types: list[str] | None = None final_logit_softcapping: float | None = None attn_logit_softcapping: float | None = None + default_theta = {"global": 1_000_000.0, "local": 10_000.0} def __post_init__(self, **kwargs): diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 15de07406bfb..1670929c8785 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -1183,7 +1183,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index a1b80a81ef80..89b926f9d4a7 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -86,6 +86,13 @@ class T5Gemma2TextConfig(Gemma3TextConfig, PreTrainedConfig): """ model_type = "t5gemma2_text" + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + use_bidirectional_attention = AttributeError() def __post_init__(self, **kwargs): @@ -973,7 +980,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) _tied_weights_keys = { "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } - _tp_plan = {"lm_head.out_proj": "colwise_gather_output"} + _tp_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/unispeech/modeling_unispeech.py b/src/transformers/models/unispeech/modeling_unispeech.py index 03103760140c..bf4fc0d108f7 100755 --- a/src/transformers/models/unispeech/modeling_unispeech.py +++ b/src/transformers/models/unispeech/modeling_unispeech.py @@ -29,8 +29,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py b/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py index 565852237d06..6cda17570834 100755 --- a/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py +++ b/src/transformers/models/unispeech_sat/modeling_unispeech_sat.py @@ -30,8 +30,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/vaultgemma/configuration_vaultgemma.py b/src/transformers/models/vaultgemma/configuration_vaultgemma.py index a60b7e8edc0c..7ff8c838d362 100644 --- a/src/transformers/models/vaultgemma/configuration_vaultgemma.py +++ b/src/transformers/models/vaultgemma/configuration_vaultgemma.py @@ -53,10 +53,25 @@ class VaultGemmaConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -64,6 +79,12 @@ class VaultGemmaConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 256000 hidden_size: int = 2304 intermediate_size: int = 9216 diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index f0a2e48d20b8..1ddba4af2210 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -467,8 +467,10 @@ def forward( @auto_docstring class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/vaultgemma/modular_vaultgemma.py b/src/transformers/models/vaultgemma/modular_vaultgemma.py index 9a4ce67c9a55..65f8b4ab3bf3 100644 --- a/src/transformers/models/vaultgemma/modular_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modular_vaultgemma.py @@ -43,6 +43,12 @@ class VaultGemmaConfig(Gemma2Config): >>> configuration = model.config ```""" + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + use_bidirectional_attention = AttributeError() diff --git a/src/transformers/models/vits/modeling_vits.py b/src/transformers/models/vits/modeling_vits.py index 8ce1411bed5e..3ab719c0921d 100644 --- a/src/transformers/models/vits/modeling_vits.py +++ b/src/transformers/models/vits/modeling_vits.py @@ -23,8 +23,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, ModelOutput diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py index a1593dfbcdca..77130e54db78 100644 --- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py @@ -30,10 +30,10 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -41,6 +41,12 @@ class VoxtralRealtimeTextConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 14336 diff --git a/src/transformers/models/wav2vec2/modeling_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_wav2vec2.py index 274a03365710..6594e813a5a9 100755 --- a/src/transformers/models/wav2vec2/modeling_wav2vec2.py +++ b/src/transformers/models/wav2vec2/modeling_wav2vec2.py @@ -26,8 +26,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer diff --git a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py index 6023d856798b..1faa3b667b73 100644 --- a/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py +++ b/src/transformers/models/wav2vec2_bert/modeling_wav2vec2_bert.py @@ -14,8 +14,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py index 710e7a64cea2..a616f62059ef 100644 --- a/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py +++ b/src/transformers/models/wav2vec2_bert/modular_wav2vec2_bert.py @@ -6,8 +6,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( diff --git a/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py b/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py index 9f35e5db42ed..aa454e20da48 100644 --- a/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py +++ b/src/transformers/models/wav2vec2_conformer/modeling_wav2vec2_conformer.py @@ -15,8 +15,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, diff --git a/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py b/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py index f02ce539d228..ce34abedbdb4 100644 --- a/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py +++ b/src/transformers/models/wav2vec2_conformer/modular_wav2vec2_conformer.py @@ -6,8 +6,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, Wav2Vec2BaseModelOutput from ...modeling_utils import PreTrainedModel diff --git a/src/transformers/models/wavlm/modeling_wavlm.py b/src/transformers/models/wavlm/modeling_wavlm.py index 18440ebf7d25..024c889ed110 100755 --- a/src/transformers/models/wavlm/modeling_wavlm.py +++ b/src/transformers/models/wavlm/modeling_wavlm.py @@ -15,8 +15,8 @@ from ... import initialization as init from ...activations import ACT2FN +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, diff --git a/src/transformers/models/wavlm/modular_wavlm.py b/src/transformers/models/wavlm/modular_wavlm.py index b3329e64913d..77a58d986ab3 100644 --- a/src/transformers/models/wavlm/modular_wavlm.py +++ b/src/transformers/models/wavlm/modular_wavlm.py @@ -5,8 +5,8 @@ import torch.nn.functional as F from ... import initialization as init +from ...distributed.fsdp import is_fsdp_managed_module from ...integrations.deepspeed import is_deepspeed_zero3_enabled -from ...integrations.fsdp import is_fsdp_managed_module from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, Wav2Vec2BaseModelOutput from ...modeling_utils import PreTrainedModel diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index 6d9f2cef1f96..6210f3c5f42b 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -53,13 +53,19 @@ class YoutuConfig(PreTrainedConfig): base_model_tp_plan = { "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } attribute_map = {} vocab_size: int = 128256 @@ -86,6 +92,7 @@ class YoutuConfig(PreTrainedConfig): rope_interleave: bool | None = True attention_bias: bool = False attention_dropout: float | int | None = 0.0 + base_model_sp_plan = None embedding_initializer_range: float | None = None def __post_init__(self, **kwargs): @@ -103,5 +110,20 @@ def __post_init__(self, **kwargs): self.head_dim = self.qk_rope_head_dim super().__post_init__(**kwargs) + def convert_rope_params_to_dict(self, **kwargs): + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs + __all__ = ["YoutuConfig"] diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index d40bef358da6..8751f452fd28 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -533,8 +533,10 @@ def forward( @auto_docstring class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index b2de3a2df0a5..f3218061670c 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -62,8 +62,9 @@ class YoutuConfig(DeepseekV3Config): base_model_tp_plan = { "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_sp_plan = None attribute_map = {} vocab_size: int = 128256 diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index d586488541df..ca2601b00fce 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -262,7 +262,10 @@ if is_torch_available(): import torch + import torch.distributed as dist + import torch.multiprocessing as mp from safetensors.torch import load_file + from torch.distributed.device_mesh import init_device_mesh from .modeling_utils import FLASH_ATTN_KERNEL_FALLBACK, PreTrainedModel @@ -316,6 +319,7 @@ def parse_int_from_env(key, default=None): _run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False) _run_training_tests = parse_flag_from_env("RUN_TRAINING_TESTS", default=True) _run_tensor_parallel_tests = parse_flag_from_env("RUN_TENSOR_PARALLEL_TESTS", default=True) +_run_fsdp_tests = parse_flag_from_env("RUN_FSDP_TESTS", default=True) def is_staging_test(test_case): @@ -382,6 +386,22 @@ def is_training_test(test_case): return pytest.mark.is_training_test()(test_case) +def is_training_distributed_test(test_case): + """ + Decorator marking a test as a training distributed test. If RUN_TRAINING_DISTRIBUTED_TESTS is set to a falsy value, those tests will be + skipped. + """ + if not _run_training_tests: + return unittest.skip(reason="test is training distributed test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_training_distributed_test()(test_case) + + def is_tensor_parallel_test(test_case): """ Decorator marking a test as a tensor parallel test. If RUN_TENSOR_PARALLEL_TESTS is set to a falsy value, those @@ -398,6 +418,22 @@ def is_tensor_parallel_test(test_case): return pytest.mark.is_tensor_parallel_test()(test_case) +def is_fsdp_test(test_case): + """ + Decorator marking a test as an FSDP test. If RUN_FSDP_TESTS is set to a falsy value, those tests will be + skipped. + """ + if not _run_fsdp_tests: + return unittest.skip(reason="test is fsdp test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_fsdp_test()(test_case) + + def slow(test_case): """ Decorator marking a test as slow. @@ -4195,6 +4231,45 @@ def read_json_file(file): # ============================================================================= +def global_wrapper(rank, func, fsdp_size, tp_size, port, func_args, func_kwargs): + def setup_dist_env(rank, world_size, port): + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + + world_size = fsdp_size * tp_size + setup_dist_env(rank, world_size, port) + + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + # NOTE(3outeille): if want to handle DataParallel, create dp_replicate dims (do not mixed with dp_shard which is for FSDP) + # NOTE(3outeille): if other parallelism is added, order matters, it should be ["pp", "ddp", "fsdp", "cp", "tp"] + # TODO(3outeille): figure out EP + # from less costly to most costly (internode to intranode) + dims, names = [fsdp_size, tp_size], ["fsdp", "tp"] + mesh = init_device_mesh("cpu", dims, mesh_dim_names=names) + + func(mesh, *func_args, **func_kwargs) + + dist.barrier() + dist.destroy_process_group() + + +def init_distributed(fsdp_size: int = 1, tp_size: int = 1): + def _init_distributed(func): + def wrapper(*args, **kwargs): + world_size = fsdp_size * tp_size + port = get_torch_dist_unique_port() + spawn_args = (func, fsdp_size, tp_size, port, args, kwargs) + mp.spawn(global_wrapper, args=spawn_args, nprocs=world_size) + + return wrapper + + return _init_distributed + + # ANSI color codes for terminal output class Colors: """ANSI color codes for terminal output formatting.""" @@ -4234,8 +4309,9 @@ class ColoredFormatter(logging.Formatter): # Loggers that should be dimmed (less important/verbose) DIMMED_LOGGERS = {"httpx", "httpcore", "urllib3", "requests"} - def __init__(self, fmt: str | None = None, datefmt: str | None = None): + def __init__(self, fmt: str | None = None, datefmt: str | None = None, rank_prefix: str = ""): super().__init__(fmt, datefmt) + self.rank_prefix = rank_prefix def format(self, record: logging.LogRecord) -> str: # Check if this logger should be dimmed @@ -4245,7 +4321,7 @@ def format(self, record: logging.LogRecord) -> str: # Dim the entire log line for httpx and similar timestamp = self.formatTime(record, self.datefmt) message = record.getMessage() - return f"{Colors.DIM}{timestamp} - {record.name} - {record.levelname:8} - {message}{Colors.RESET}" + return f"{Colors.DIM}{timestamp} - {record.name} - {record.levelname:8} - {self.rank_prefix}{message}{Colors.RESET}" # Get color for this level color = self.LEVEL_COLORS.get(record.levelno, Colors.RESET) @@ -4263,7 +4339,7 @@ def format(self, record: logging.LogRecord) -> str: # Get message message = record.getMessage() - return f"{colored_time} - {colored_name} - {colored_levelname} - {message}" + return f"{colored_time} - {colored_name} - {colored_levelname} - {self.rank_prefix}{message}" _warn_once_logged: set[str] = set() @@ -4274,26 +4350,34 @@ def init_test_logger() -> logging.Logger: Uses a named logger instead of root logger to avoid conflicts with pytest-xdist parallel execution. Uses stderr instead of stdout to avoid deadlocks with pytest-xdist output capture. + Automatically includes rank in log format when distributed is initialized. """ logger = logging.getLogger("transformers.training_test") logger.setLevel(logging.INFO) - # Only add handler if not already present (avoid duplicate handlers on repeated calls) - if not logger.handlers: - # Use stderr instead of stdout - pytest-xdist captures stdout which can cause deadlocks - ch = logging.StreamHandler(sys.stderr) - ch.setLevel(logging.INFO) + # Clear existing handlers to update format (e.g., when dist becomes initialized) + logger.handlers.clear() - # Use colored formatter if terminal supports it, plain otherwise - if sys.stderr.isatty(): - formatter = ColoredFormatter(datefmt="%Y-%m-%d %H:%M:%S") - else: - formatter = logging.Formatter( - "%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S" - ) + # Use stderr instead of stdout - pytest-xdist captures stdout which can cause deadlocks + ch = logging.StreamHandler(sys.stderr) + ch.setLevel(logging.INFO) + + # Build format string - include rank if distributed is initialized + rank_prefix = "" + if is_torch_available() and dist.is_initialized(): + rank = dist.get_rank() + rank_prefix = f"[rank{rank}] " + + # Use colored formatter if terminal supports it, plain otherwise + if sys.stderr.isatty(): + formatter = ColoredFormatter(datefmt="%Y-%m-%d %H:%M:%S", rank_prefix=rank_prefix) + else: + formatter = logging.Formatter( + f"%(asctime)s - %(name)s - %(levelname)s - {rank_prefix}%(message)s", datefmt="%Y-%m-%d %H:%M:%S" + ) - ch.setFormatter(formatter) - logger.addHandler(ch) + ch.setFormatter(formatter) + logger.addHandler(ch) logger.propagate = False # Don't propagate to root logger to avoid duplicate output return logger diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index ad66f1a23b0c..2eca1e64dda1 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -55,6 +55,7 @@ from .configuration_utils import PreTrainedConfig from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator from .debug_utils import DebugOption, DebugUnderflowOverflow +from .distributed.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .feature_extraction_sequence_utils import SequenceFeatureExtractor from .feature_extraction_utils import FeatureExtractionMixin from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend @@ -66,7 +67,6 @@ is_deepspeed_available, propagate_args_to_deepspeed, ) -from .integrations.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .integrations.liger import apply_liger_kernel from .integrations.neftune import activate_neftune, deactivate_neftune from .integrations.peft import MIN_PEFT_VERSION @@ -2414,9 +2414,10 @@ def get_cp_size(self) -> int: def get_tp_size(self) -> int: """Get the tensor parallel size from either the model or DeepSpeed config.""" - # 1. Check model.tp_size first - if (model_tp := getattr(self.model, "_tp_size", None)) is not None: - return model_tp + # TODO: adapt it cleaner with distributed config once distributed api is stable + dc = getattr(getattr(self.model, "config", None), "distributed_config", None) + if dc is not None and dc.tp_size is not None: + return dc.tp_size # 2. Fall back to DeepSpeed config if enabled if self.is_deepspeed_enabled and (deepspeed_config := getattr(self.args, "hf_deepspeed_config", None)): diff --git a/src/transformers/trainer_seq2seq.py b/src/transformers/trainer_seq2seq.py index ada588adbd21..c3907ed2556c 100644 --- a/src/transformers/trainer_seq2seq.py +++ b/src/transformers/trainer_seq2seq.py @@ -22,9 +22,9 @@ from torch import nn from torch.utils.data import Dataset +from .distributed.fsdp import is_fsdp_managed_module from .generation.configuration_utils import GenerationConfig from .integrations.deepspeed import is_deepspeed_zero3_enabled -from .integrations.fsdp import is_fsdp_managed_module from .trainer import Trainer from .utils import is_datasets_available, logging diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py index 6b94a520d4f2..9df6b77447ae 100644 --- a/tests/causal_lm_tester.py +++ b/tests/causal_lm_tester.py @@ -30,6 +30,7 @@ ) from .test_configuration_common import ConfigTester +from .test_fsdp_mixin import FSDPTesterMixin from .test_modeling_common import ( GenerationTesterMixin, ModelTesterMixin, @@ -271,7 +272,12 @@ def prepare_config_and_inputs_for_common(self): @require_torch class CausalLMModelTest( - ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TrainingTesterMixin, TensorParallelTesterMixin + ModelTesterMixin, + GenerationTesterMixin, + PipelineTesterMixin, + TrainingTesterMixin, + TensorParallelTesterMixin, + FSDPTesterMixin, ): model_tester_class = None all_model_classes = None diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 91770b683e45..e81aa342d635 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -11,56 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import math import warnings -from types import SimpleNamespace - -import torch from transformers import AutoModelForCausalLM -from transformers.integrations.tensor_parallel import ( - ColwiseParallel, - EmbeddingParallel, - GroupedGemmParallel, - PackedColwiseParallel, - PackedRowwiseParallel, - RowwiseParallel, - get_packed_weights, - repack_weights, -) from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test -@is_tensor_parallel_test -class TestTensorParallelUtils(TestCasePlus): - def test_packed_unpacked_conversion(self): - WORLD_SIZE = 2 - PACKED_BLOCK_SIZE = 800 - SHARDING_DIM = 2 - NUM_BLOCKS = 2 - - original_packed_weights = torch.randn(4, 512, 2 * PACKED_BLOCK_SIZE) - original_packed_weights.get_dtype = lambda: "F32" # get_packed_weights expects PySlice object - empty_param = torch.empty(4, 512, 2 * PACKED_BLOCK_SIZE) - - class MockDeviceMesh: - def size(self): - return WORLD_SIZE - - mock_mesh = ( - MockDeviceMesh() - ) # get_packed_weights only calls `.size()`, do this to avoid doing actual distributed run - - packed_weights_0 = get_packed_weights(original_packed_weights, empty_param, mock_mesh, 0, SHARDING_DIM) - packed_weights_1 = get_packed_weights(original_packed_weights, empty_param, mock_mesh, 1, SHARDING_DIM) - - # simulate all gather of sharded weights - packed_weights = torch.cat([packed_weights_0, packed_weights_1], dim=SHARDING_DIM) - unpacked_weights = repack_weights(packed_weights, SHARDING_DIM, WORLD_SIZE, NUM_BLOCKS) - - assert torch.allclose(unpacked_weights, original_packed_weights) - - @is_tensor_parallel_test class TestTensorParallelProperties(TestCasePlus): def test_tp_plan_property_setter_getter(self): @@ -90,18 +46,6 @@ def test_tp_plan_property_setter_getter(self): } self.assertEqual(model.tp_plan, expected_plan) - def test_tp_plan_validation_invalid_style(self): - """Test that invalid parallel styles are rejected.""" - model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto") - - # Test invalid parallel style - with self.assertRaises(ValueError) as context: - model.tp_plan = {"layers.*.self_attn.q_proj": "invalid_style"} - - self.assertIn("Unsupported tensor parallel style 'invalid_style'", str(context.exception)) - self.assertIn("Supported styles are", str(context.exception)) - def test_tp_plan_validation_nonexistent_layer_warning(self): """Test that warnings are issued for non-existent layer patterns.""" @@ -165,256 +109,3 @@ def test_tp_plan_none_handling(self): # Test setting a plan after None model.tp_plan = {"model.layers.*.self_attn.q_proj": "colwise"} self.assertEqual(model.tp_plan, {"model.layers.*.self_attn.q_proj": "colwise"}) - - -@is_tensor_parallel_test -class TestTensorParallelLayer(TestCasePlus): - class MockDeviceMesh: - def __init__(self, world_size, rank): - self.world_size = world_size - self.rank = rank - self.shape = (world_size,) - - def size(self): - return self.world_size - - def get_local_rank(self): - return self.rank - - def test_colwise_get_expected_sharded_shape(self): - world_size = 3 - size = 10 # not divisible by world_size to test edge case - empty_param_2d = torch.empty(size, 32) - empty_param_1d = torch.empty((size,)) - step = math.ceil(size / world_size) - - for rank in range(world_size): - for empty_param in [empty_param_2d, empty_param_1d]: - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - layer = ColwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param) - - begin = rank * step - end = min(begin + step, size) - ground_truth = (end - begin,) + empty_param.shape[1:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - def test_rowwise_get_expected_sharded_shape(self): - world_size = 3 - size = 10 # not divisible by world_size to test edge case - empty_param_2d = torch.empty(32, size) - empty_param_1d = torch.empty((size,)) - step = math.ceil(size / world_size) - - for rank in range(world_size): - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - - # 2D: shards on dim -1 (input features) - layer = RowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param_2d) - begin = rank * step - end = min(begin + step, size) - ground_truth = empty_param_2d.shape[:-1] + (end - begin,) - expected_shape = layer.get_expected_sharded_shape(empty_param_2d.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - # 1D bias: NOT sharded - layer = RowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param_1d) - self.assertEqual(layer.get_expected_sharded_shape(empty_param_1d.shape), empty_param_1d.shape) - - def test_embedding_get_expected_sharded_shape(self): - world_size = 3 - size = 10 # not divisible by world_size to test edge case; same size on both dims so step applies to both - empty_param = torch.empty(size, size) - step = math.ceil(size / world_size) - - for rank in range(world_size): - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - begin = rank * step - end = min(begin + step, size) - - # embedding_dim_sharding=0: shards dim 0 (vocab) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=rank, empty_param=empty_param, embedding_dim_sharding=0 - ) - ground_truth = (end - begin,) + empty_param.shape[1:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - # embedding_dim_sharding=1: shards dim 1 (embedding dim) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=rank, empty_param=empty_param, embedding_dim_sharding=1 - ) - ground_truth = empty_param.shape[:1] + (end - begin,) + empty_param.shape[2:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - def test_grouped_gemm_get_expected_sharded_shape(self): - world_size = 3 - size = 9 # must be divisible by world_size (GroupedGemm requires it) - empty_param = torch.empty(size, 16, 32) - step = math.ceil(size / world_size) - - for rank in range(world_size): - device_mesh = self.MockDeviceMesh(world_size=world_size, rank=rank) - layer = GroupedGemmParallel(device_mesh=device_mesh, rank=rank, empty_param=empty_param) - begin = rank * step - end = min(begin + step, size) - ground_truth = (end - begin,) + empty_param.shape[1:] - expected_shape = layer.get_expected_sharded_shape(empty_param.shape) - self.assertEqual( - expected_shape, ground_truth, f"Rank {rank} expected shape {ground_truth} but got {expected_shape}" - ) - - def test_colwise_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - # gather_output=False (default): out_features is updated - module = torch.nn.Linear(32, 16) - layer = ColwiseParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32)) - layer.update_module_attributes(module) - self.assertEqual(module.out_features, 4) - - # gather_output=True: out_features is NOT updated - module = torch.nn.Linear(32, 16) - layer = ColwiseParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32), gather_output=True) - layer.update_module_attributes(module) - self.assertEqual(module.out_features, 16) - - def test_rowwise_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - module = torch.nn.Linear(32, 16) - layer = RowwiseParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32)) - layer.update_module_attributes(module) - self.assertEqual(module.in_features, 8) - - def test_embedding_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - # embedding_dim_sharding=0: num_embeddings is updated - module = torch.nn.Embedding(32, 16) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=0, empty_param=torch.empty(32, 16), embedding_dim_sharding=0 - ) - layer.update_module_attributes(module) - self.assertEqual(module.num_embeddings, 8) - self.assertEqual(module.embedding_dim, 16) - - # embedding_dim_sharding=1: embedding_dim is updated - module = torch.nn.Embedding(32, 16) - layer = EmbeddingParallel( - device_mesh=device_mesh, rank=0, empty_param=torch.empty(32, 16), embedding_dim_sharding=1 - ) - layer.update_module_attributes(module) - self.assertEqual(module.num_embeddings, 32) - self.assertEqual(module.embedding_dim, 4) - - def test_grouped_gemm_update_module_attributes(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - - # There is no torch module with num_experts attribute, it is more at the Transformers level, - # so just use a SimpleNamespace to test that the attribute is updated correctly. - module = SimpleNamespace(num_experts=8) - layer = GroupedGemmParallel(device_mesh=device_mesh, rank=0, empty_param=torch.empty(8, 16, 32)) - layer.update_module_attributes(module) - self.assertEqual(module.num_experts, 2) - - def test_update_module_attributes_missing_attribute(self): - device_mesh = self.MockDeviceMesh(world_size=4, rank=0) - module = SimpleNamespace(random_attr=123) - for cls in [ColwiseParallel, RowwiseParallel, GroupedGemmParallel]: - layer = cls(device_mesh=device_mesh, rank=0, empty_param=torch.empty(16, 32)) - layer.update_module_attributes(module) - - self.assertEqual( - module.__dict__, - {"random_attr": 123}, - "update_module_attributes should not modify attributes that don't exist", - ) - - def test_shard_tensor_shape_consistency(self): - """ - Test that shard_tensor returns tensors of the expected shape for different parallel styles and ranks. - """ - WORLD_SIZE = 4 - cases = [ - (ColwiseParallel, (16, 32), {}), - (ColwiseParallel, (16, 32), {"gather_output": True}), - (ColwiseParallel, (16,), {}), - (RowwiseParallel, (16, 32), {}), - (RowwiseParallel, (32,), {}), - (EmbeddingParallel, (32, 16), {"embedding_dim_sharding": 0}), - (EmbeddingParallel, (32, 16), {"embedding_dim_sharding": 1}), - ] - for cls, shape, kwargs in cases: - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = cls(device_mesh=device_mesh, rank=rank, empty_param=torch.empty(*shape), **kwargs) - - full_tensor = torch.randn(*shape) - sharded = layer.shard_tensor(full_tensor) - expected = layer.get_expected_sharded_shape(shape) - - self.assertEqual(tuple(sharded.shape), expected, f"{cls.__name__} rank={rank} shape={shape}") - - def test_packed_colwise_shard_tensor(self): - WORLD_SIZE = 2 - # 3D empty_param - empty = torch.empty(2, 16, 64) - - # Packed vs unpacked path is determined by checking the following: - # input.dim() == get_expected_sharded_shape(empty_param).dim() - - # Packed - full_packed = torch.randn(2, 16, 64) - full_packed.get_dtype = lambda: "F32" - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedColwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_packed) - expected_shape = (2, 8, 64) # last dim is packed size, middle dim is sharded - self.assertEqual(sharded.shape, expected_shape) - - # Unpacked - full_unpacked = torch.randn(16, 64) - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedColwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_unpacked) - expected_shape = (8, 64) # last dim is not packed, so just sharded - self.assertEqual(sharded.shape, expected_shape) - - def test_packed_rowwise_shard_tensor(self): - WORLD_SIZE = 2 - # empty_param last dim = 64 signals the packed size (2 * 32) - empty = torch.empty(16, 64) - - # Packed vs unpacked path is determined by checking the following: - # input.shape[-1] < empty_param.shape[-1] - - # Packed - full_packed = torch.randn(16, 64) - full_packed.get_dtype = lambda: "F32" - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedRowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_packed) - expected_shape = (16, 32) # last dim is packed size, sharded - self.assertEqual(sharded.shape, expected_shape) - - # Unpacked - full_unpacked = torch.randn(16, 32) - for rank in range(WORLD_SIZE): - device_mesh = self.MockDeviceMesh(world_size=WORLD_SIZE, rank=rank) - layer = PackedRowwiseParallel(device_mesh=device_mesh, rank=rank, empty_param=empty) - sharded = layer.shard_tensor(full_unpacked) - expected_shape = (16, 16) # last dim is not packed, so just sharded - self.assertEqual(sharded.shape, expected_shape) diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py new file mode 100644 index 000000000000..3cf5b4a914db --- /dev/null +++ b/tests/test_distributed_config.py @@ -0,0 +1,93 @@ +import json +import tempfile + +from transformers.distributed import DistributedConfig + + +class TestDistributedConfig: + def test_2d_parallelism(self): + dc = DistributedConfig(tp_size=2, fsdp_size=2) + assert dc.tp_size == 2 + assert dc.fsdp_size == 2 + assert dc.tp_plan is None + assert dc.fsdp_plan is None + + def test_tp_only_defaults_fsdp_to_1(self): + dc = DistributedConfig(tp_size=4) + assert dc.tp_size == 4 + assert dc.fsdp_size == 1 + assert dc.tp_plan is None + + def test_fsdp_only_defaults_tp_to_1(self): + dc = DistributedConfig(fsdp_size=4) + assert dc.tp_size == 1 + assert dc.fsdp_size == 4 + assert dc.fsdp_plan is None + assert dc.tp_plan is None + + def test_empty_config(self): + dc = DistributedConfig() + assert dc.tp_size is None + assert dc.fsdp_size is None + assert dc.tp_plan is None + assert dc.fsdp_plan is None + + def test_from_dict(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) + assert dc.tp_size == 2 + assert dc.fsdp_size == 4 + assert dc.tp_plan is None + + def test_from_dict_ignores_unknown_keys(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "unknown_key": 42}) + assert dc.tp_size == 2 + assert not hasattr(dc, "unknown_key") + + def test_from_dict_kwargs_override(self): + dc = DistributedConfig.from_dict({"tp_size": 2}, tp_size=8) + assert dc.tp_size == 8 + + def test_to_dict(self): + dc = DistributedConfig(tp_size=2, fsdp_size=4) + d = dc.to_dict() + assert d == { + "tp_size": 2, + "tp_plan": None, + "enable_sequence_parallel": False, + "enable_expert_parallel": False, + "fsdp_size": 4, + "fsdp_plan": None, + } + + def test_to_dict_is_a_copy(self): + dc = DistributedConfig(tp_plan={"layer": "colwise"}) + d = dc.to_dict() + d["tp_plan"]["layer"] = "rowwise" + assert dc.tp_plan["layer"] == "colwise" + + def test_to_json_string(self): + dc = DistributedConfig(tp_size=2, fsdp_size=2) + s = dc.to_json_string() + parsed = json.loads(s) + assert parsed["tp_size"] == 2 + assert parsed["fsdp_size"] == 2 + + def test_to_json_file(self): + dc = DistributedConfig(tp_size=4) + with tempfile.NamedTemporaryFile(mode="r", suffix=".json", delete=False) as f: + dc.to_json_file(f.name) + f.seek(0) + parsed = json.load(f) + assert parsed["tp_size"] == 4 + assert parsed["tp_plan"] is None + + def test_roundtrip_dict(self): + original = DistributedConfig(tp_size=2, fsdp_size=4) + restored = DistributedConfig.from_dict(original.to_dict()) + assert original == restored + + def test_repr(self): + dc = DistributedConfig(tp_size=2) + r = repr(dc) + assert "DistributedConfig" in r + assert '"tp_size": 2' in r diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py new file mode 100644 index 000000000000..88734b501dfe --- /dev/null +++ b/tests/test_fsdp_mixin.py @@ -0,0 +1,784 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FSDP tester mixin for model tests.""" + +import json +import logging +import os +import socket +import sys +import tempfile +import time +import traceback +from abc import ABC, abstractmethod + +from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, is_torch_available +from transformers.testing_utils import ( + backend_device_count, + backend_empty_cache, + backend_torch_accelerator_module, + init_test_logger, + is_fsdp_test, + require_fsdp, +) +from transformers.trainer_utils import set_seed + + +logger = logging.getLogger("transformers.training_test") + + +if is_torch_available(): + import torch + import torch.distributed as dist + import torch.multiprocessing as mp + from torch.nn.parallel import DistributedDataParallel as DDP + + from transformers.distributed import DistributedConfig + from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard + from transformers.distributed.utils import ( + gather_full_state_dict, + load_optimizer_distributed, + save_optimizer_distributed, + ) + + +# ============================================================================= +# Constants +# ============================================================================= + +BATCH_SIZE = 2 +SEQ_LEN = 64 +NUM_STEPS = 20 +LR = 3e-4 +SEED = 42 +FSDP_TOP_MODEL_NAMES = { + # FSDP coverage is gated on models declaring `base_model_fsdp_plan` (config) + + # class-level `_fsdp_plan` (head class). Listed here are models whose test + # class extends `CausalLMModelTest` (so they pick up the FSDP mixin) and + # which use the standard embed_tokens / layers.* / norm naming. + # Dense + "llama", + "mistral", + "qwen3", + "phi", + "olmo3", + "gemma2", + # MoE + "mixtral", + "qwen3_moe", + "qwen2_moe", + "qwen3_5_moe", + "deepseek_v2", + "gpt_oss", + "glm4_moe_lite", +} + + +# ============================================================================= +# Distributed helpers (top-level for pickling by mp.spawn) +# ============================================================================= + + +def _get_distributed_device_type(): + device_type = torch._C._get_accelerator().type + return "cpu" if device_type == "mps" else device_type + + +def _get_distributed_backend(): + backend_map = {"cpu": "gloo", "cuda": "nccl", "xpu": "xccl", "hpu": "hccl"} + return backend_map.get(_get_distributed_device_type(), "gloo") + + +def _get_rank_device(rank): + device_type = _get_distributed_device_type() + if device_type == "cpu": + return torch.device("cpu") + return torch.device(device_type, rank) + + +def _set_rank_device(rank): + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if accelerator_module is not None and hasattr(accelerator_module, "set_device"): + accelerator_module.set_device(rank) + + +def _get_accelerator_rng_state(): + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if accelerator_module is None or not hasattr(accelerator_module, "get_rng_state"): + return None + return accelerator_module.get_rng_state() + + +def _set_accelerator_rng_state(rng_state): + accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type()) + if rng_state is not None and accelerator_module is not None and hasattr(accelerator_module, "set_rng_state"): + accelerator_module.set_rng_state(rng_state) + + +def _get_available_fsdp_workers(): + if _get_distributed_device_type() == "cpu": + return os.cpu_count() or 1 + return backend_device_count(_get_distributed_device_type()) + + +def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_size, port, results_file): + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + + _set_determinism(SEED) + + dist.init_process_group(backend=_get_distributed_backend(), rank=rank, world_size=world_size) + _set_rank_device(rank) + + if rank == 0: + start_time = time.perf_counter() + print(f"[FSDP] Starting test: {test_name}", flush=True) + + error = None + try: + func(rank, *func_args, **func_kwargs) + except Exception as e: + error = f"{type(e).__name__}: {e}\n{traceback.format_exc()}" + + error_flag = torch.tensor([1 if error else 0], device=_get_rank_device(rank)) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX) + any_failed = error_flag.item() > 0 + + if rank == 0: + elapsed = time.perf_counter() - start_time + status = "FAIL" if any_failed else "PASS" + output_stream = sys.stderr if any_failed else sys.stdout + print(f"[FSDP] {status} test: {test_name} ({elapsed:.1f}s)", file=output_stream, flush=True) + with open(results_file, "w") as f: + json.dump({"error": error or ("Failed on another rank" if any_failed else None)}, f) + + backend_empty_cache(_get_distributed_device_type()) + dist.barrier() + dist.destroy_process_group() + + +def _set_determinism(seed): + torch.use_deterministic_algorithms(True) + if _get_distributed_device_type() == "cuda" and torch.cuda.is_available(): + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + set_seed(seed) + + +# ============================================================================= +# Training & comparison helpers (top-level for pickling) +# ============================================================================= + + +def _build_repeated_training_batches(config, device, num_steps): + """Create one deterministic batch and reuse it across steps.""" + generator = torch.Generator(device=device) + generator.manual_seed(SEED) + input_ids = torch.randint(0, config.vocab_size, (BATCH_SIZE, SEQ_LEN), device=device, generator=generator) + labels = input_ids.clone() + return [(input_ids, labels)] * num_steps + + +def _create_shared_tmpdir(rank): + if rank == 0: + tmpdir_obj = tempfile.TemporaryDirectory() + tmpdir = tmpdir_obj.name + tmpdir_list = [tmpdir] + else: + tmpdir_obj = None + tmpdir_list = [None] + dist.broadcast_object_list(tmpdir_list, src=0) + return tmpdir_list[0], tmpdir_obj + + +def _gather_ddp_state_dict(model): + # Only rank 0 returns data to match gather_full_state_dict semantics, so the + # downstream DDP-vs-FSDP comparison runs once on rank 0 instead of N times. + if dist.get_rank() != 0: + return {} + return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()} + + +def _resolve_fsdp_plan_paths(model): + """Expand model._fsdp_plan into (paths, strategy) entries. + + Wildcard keys are expanded via ``replace_layer_number_by_wildcard``. When + weights are tied, the standalone embed_tokens entry is skipped and any + ``"lm_head"`` keep entry is rewritten to the tied source path (the keep + group will wrap the shared parameter once). + """ + plan = model._fsdp_plan + + input_embed = model.get_input_embeddings() + output_embed = model.get_output_embeddings() + weights_tied = ( + input_embed is not None + and output_embed is not None + and hasattr(input_embed, "weight") + and hasattr(output_embed, "weight") + and input_embed.weight is output_embed.weight + ) + tied_source = None + if weights_tied: + for name, mod in model.named_modules(): + if mod is input_embed: + tied_source = name + break + + name_to_module = dict(model.named_modules()) + entries: list[tuple[list[str], str]] = [] + for key, strategy in plan.items(): + if weights_tied and key == tied_source: + continue + if weights_tied and key == "lm_head" and strategy == "keep_full_weight": + entries.append(([tied_source], strategy)) + continue + + if key in name_to_module: + entries.append(([key], strategy)) + continue + matched = [name for name in name_to_module if replace_layer_number_by_wildcard(name) == key] + if matched: + entries.append((matched, strategy)) + return entries + + +def _build_manual_fsdp_plan(config, device, policy_options=None): + """Build a manual FSDP2 plan by expanding model._fsdp_plan.""" + policy_options = policy_options or [] + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(device) + + module_plan: dict[str, list[str]] = {} + for paths, strategy in _resolve_fsdp_plan_paths(model): + for path in paths: + module_plan[path] = [strategy, *policy_options] + + del model + return {"modules": module_plan} + + +def _save_init_pretrained(rank, config, dtype): + """Save a deterministic initial model to a shared tmpdir for from_pretrained loading.""" + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + if rank == 0: + set_seed(SEED) + model = AutoModelForCausalLM.from_config(config).to(dtype) + model.save_pretrained(tmpdir) + del model + dist.barrier() + return tmpdir, tmpdir_obj + + +def _save_training_state(model, optimizer, training_state_dir): + """Save optimizer (canonical DCP path) plus per-rank RNG for resume.""" + save_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim")) + rng = {"cpu": torch.get_rng_state()} + accel = _get_accelerator_rng_state() + if accel is not None: + rng["accel"] = accel + torch.save(rng, os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt")) + + +def _load_training_state(model, optimizer, training_state_dir): + """Inverse of `_save_training_state`.""" + load_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim")) + rng = torch.load(os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt"), weights_only=False) + torch.set_rng_state(rng["cpu"]) + if "accel" in rng: + _set_accelerator_rng_state(rng["accel"]) + + +def train_ddp(rank, batches, lr, device, dtype, init_model_dir): + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype).to(device) + # MoE/conditional-routing variants) may not use all params on + # every step, and DDP would otherwise fail. Specifying find_unused_parameters=True allows running backward on a subgraph of the model. + ddp_kwargs = {"find_unused_parameters": True} + if device.type != "cpu": + ddp_kwargs["device_ids"] = [rank] + ddp_model = DDP(model, **ddp_kwargs) + ddp_model.train() + optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr) + + losses, grad_norms = [], [] + for input_ids, labels in batches: + optimizer.zero_grad() + output = ddp_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=float("inf")) + optimizer.step() + + losses.append(loss.detach().item()) + grad_norms.append(grad_norm) + + state_dict = _gather_ddp_state_dict(ddp_model) + + del optimizer, ddp_model, model + backend_empty_cache(_get_distributed_device_type()) + dist.barrier() + + return losses, grad_norms, state_dict + + +def train_fsdp2( + rank, + batches, + lr, + dtype, + init_model_dir, + checkpoint_step, + fsdp_plan, +): + # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save + _set_determinism(SEED) + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size(), fsdp_plan=fsdp_plan) + pre_ckpt_model = AutoModelForCausalLM.from_pretrained( + init_model_dir, torch_dtype=dtype, distributed_config=distributed_config + ) + pre_ckpt_model.train() + pre_ckpt_optimizer = torch.optim.Adam(pre_ckpt_model.parameters(), lr=lr) + + pre_ckpt_losses, pre_ckpt_grad_norms = [], [] + for step in range(0, checkpoint_step): + input_ids, labels = batches[step] + pre_ckpt_optimizer.zero_grad() + output = pre_ckpt_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(pre_ckpt_model.parameters(), max_norm=float("inf")) + pre_ckpt_optimizer.step() + + pre_ckpt_losses.append(loss.detach().item()) + pre_ckpt_grad_norms.append(grad_norm) + + # -- Phase 2: Save checkpoint, then load into a fresh model + # tmpdir/ + # model/ <- HF safetensors via save_pretrained (DCP + consolidation) + # training_state/ <- distcp (optimizer + RNG) + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + try: + model_dir = os.path.join(tmpdir, "model") + training_state_dir = os.path.join(tmpdir, "training_state") + + pre_ckpt_model.save_pretrained(model_dir, is_main_process=(rank == 0)) + _save_training_state(pre_ckpt_model, pre_ckpt_optimizer, training_state_dir) + dist.barrier() + + # Intentionally scramble RNG to prove checkpoint restore works + _set_determinism(SEED + 1234) + resumed_model = AutoModelForCausalLM.from_pretrained( + model_dir, torch_dtype=dtype, distributed_config=distributed_config + ) + resumed_model.train() + resumed_optimizer = torch.optim.Adam(resumed_model.parameters(), lr=lr) + + _load_training_state(resumed_model, resumed_optimizer, training_state_dir) + dist.barrier() + finally: + if rank == 0: + tmpdir_obj.cleanup() + + # -- Phase 3: Post-checkpoint run -- continue training the remaining steps from the resumed model + post_ckpt_losses, post_ckpt_grad_norms = [], [] + for step in range(checkpoint_step, len(batches)): + input_ids, labels = batches[step] + resumed_optimizer.zero_grad() + output = resumed_model(input_ids=input_ids, labels=labels, use_cache=False) + loss = output.loss + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(resumed_model.parameters(), max_norm=float("inf")) + resumed_optimizer.step() + + post_ckpt_losses.append(loss.detach().item()) + post_ckpt_grad_norms.append(grad_norm) + + combined_losses = pre_ckpt_losses + post_ckpt_losses + combined_grad_norms = pre_ckpt_grad_norms + post_ckpt_grad_norms + combined_state_dict = gather_full_state_dict(resumed_model) + + return combined_losses, combined_grad_norms, combined_state_dict + + +# ============================================================================= +# Distributed test implementations (top-level for pickling by mp.spawn) +# ============================================================================= +def _test_fsdp2_save_load_impl(rank, config_class, config_dict): + """Train FSDP2 model, save via save_pretrained, load via from_pretrained, compare state dicts.""" + init_test_logger() + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + + batches = _build_repeated_training_batches(config, device, 3) + + distributed_config = DistributedConfig(fsdp_size=dist.get_world_size()) + + init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) + try: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained(init_tmpdir, distributed_config=distributed_config) + dist.barrier() + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + model.train() + optimizer = torch.optim.Adam(model.parameters(), lr=LR) + + for input_ids, labels in batches: + optimizer.zero_grad() + output = model(input_ids=input_ids, labels=labels, use_cache=False) + output.loss.backward() + optimizer.step() + + state_dict_before = gather_full_state_dict(model) + + tmpdir, tmpdir_obj = _create_shared_tmpdir(rank) + try: + model.save_pretrained(tmpdir, is_main_process=(rank == 0)) + dist.barrier() + + new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config) + dist.barrier() + finally: + if rank == 0: + tmpdir_obj.cleanup() + + state_dict_after = gather_full_state_dict(new_model) + + for key in state_dict_before: + assert key in state_dict_after, f"Key {key} missing after load" + torch.testing.assert_close( + state_dict_before[key], + state_dict_after[key], + rtol=0, + atol=0, + msg=f"Weight mismatch for {key} after save/load", + ) + + if rank == 0: + logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.") + + +def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings): + """ + Verify that apply_fully_shard_data_parallel(fsdp_plan=None) wraps exactly the right modules. + + Expected FSDP targets: + UNTIED TIED + ────── ──── + 1. embed_tokens (reshard=True) 1. (skip — embed goes to step 3) + 2. layers[i] (reshard=True) 2. layers[i] (reshard=True) + 3. [norm, lm_head] (reshard=False) 3. [norm, embed_tokens] (reshard=False) + 4. root 4. root + """ + init_test_logger() + + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32) + try: + _set_determinism(SEED) + model = AutoModelForCausalLM.from_pretrained( + init_tmpdir, distributed_config=DistributedConfig(fsdp_size=dist.get_world_size()) + ) + dist.barrier() + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + + # Expected FSDP targets come from model._fsdp_plan: every resolved path gets a + # fully_shard call (keep_full_weight entries are bundled into one group, but each + # member still appears as an FSDP-wrapped module in named_modules), plus the root. + expected_targets = {""} + for paths, _strategy in _resolve_fsdp_plan_paths(model): + expected_targets.update(paths) + + actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} + + if rank == 0: + logger.debug(f" Weights tied: {config.tie_word_embeddings}") + logger.debug(f" Expected FSDP targets: {sorted(expected_targets)}") + logger.debug(f" Actual FSDP targets: {sorted(actual_targets)}") + + missing = expected_targets - actual_targets + extra = actual_targets - expected_targets + assert not missing and not extra, ( + f"FSDP target mismatch.\n" + f" Missing (expected but not wrapped): {sorted(missing)}\n" + f" Extra (wrapped but not expected): {sorted(extra)}" + ) + + if rank == 0: + logger.debug(f" FSDP sharding structure OK ({len(actual_targets)} targets)") + + +def _test_fsdp2_plan_vs_ddp_impl( + rank, config_class, config_dict, tie_word_embeddings, plan_mode, policy_options=None, dtype=None +): + """Validate DDP-vs-FSDP2 trace matching for either auto or manual plan mode.""" + init_test_logger() + + if dtype is None: + dtype = torch.float32 + + policy_options = policy_options or [] + assert "mixed_precision" not in policy_options, ( + "Use the mixed-precision specific tests when enabling mixed_precision policy." + ) + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + config.tie_word_embeddings = tie_word_embeddings + + if plan_mode == "auto": + fsdp_plan = { + "cpu_offload": "cpu_offload" in policy_options, + "mixed_precision": "mixed_precision" in policy_options, + } + test_label = f"FSDP2(auto{'+policies' if policy_options else ''})" + elif plan_mode == "manual": + fsdp_plan = _build_manual_fsdp_plan(config, device, policy_options=policy_options) + test_label = f"FSDP2(manual{'+policies' if policy_options else ''})" + else: + raise ValueError(f"Unsupported plan_mode '{plan_mode}'. Expected 'auto' or 'manual'.") + + checkpoint_step = NUM_STEPS // 2 + init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype) + batches = _build_repeated_training_batches(config, device, NUM_STEPS) + try: + ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir) + + fsdp_losses, fsdp_grad_norms, fsdp_state_dict = train_fsdp2( + rank, + batches, + LR, + dtype, + init_model_dir=init_model_dir, + checkpoint_step=checkpoint_step, + fsdp_plan=fsdp_plan, + ) + finally: + if rank == 0 and init_tmpdir_obj is not None: + init_tmpdir_obj.cleanup() + + for step in range(len(ddp_losses)): + torch.testing.assert_close( + torch.tensor(ddp_losses[step]), + torch.tensor(fsdp_losses[step]), + rtol=1e-5, + atol=1e-5, + msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}", + ) + torch.testing.assert_close( + torch.tensor(ddp_grad_norms[step]), + torch.tensor(fsdp_grad_norms[step]), + rtol=1e-5, + atol=1e-5, + msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}", + ) + + for key in ddp_state_dict: + assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict" + torch.testing.assert_close( + ddp_state_dict[key], + fsdp_state_dict[key], + rtol=1e-5, + atol=1e-5, + msg=f"Weight mismatch for {key}: DDP vs {test_label}", + ) + + if rank == 0: + logger.debug(f"DDP and {test_label} comparison checks passed.") + + +# ============================================================================= +# Mixin class +# ============================================================================= + + +class FSDPTesterMixin(ABC): + fsdp_nproc_per_node: int = 2 + skip_fsdp_tests: bool = False + + @property + @abstractmethod + def model_tester(self): + """The model tester instance (e.g., CausalLMModelTester).""" + ... + + def _skip_if_fsdp_disabled(self): + if self.skip_fsdp_tests: + self.skipTest("FSDP tests disabled for this model (skip_fsdp_tests=True)") + + def _skip_if_insufficient_devices(self): + self._skip_if_fsdp_disabled() + available_workers = _get_available_fsdp_workers() + if available_workers < self.fsdp_nproc_per_node: + self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}") + + def _get_fsdp_model_name(self): + module_parts = self.__class__.__module__.split(".") + if len(module_parts) >= 3 and module_parts[0] == "tests" and module_parts[1] == "models": + return module_parts[2] + return None + + def _skip_if_fsdp_model_not_selected(self): + model_name = self._get_fsdp_model_name() + if model_name not in FSDP_TOP_MODEL_NAMES: + model_label = model_name or self.__class__.__module__ + self.skipTest( + "FSDP mixin coverage is currently limited to the top-10 dense and top-10 MoE model suites " + f"(skipping {model_label})." + ) + + def _create_model_on_meta(self, config): + """Instantiate a model on the meta device (no memory allocated).""" + auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM] + for auto_cls in auto_classes: + try: + with torch.device("meta"): + return auto_cls.from_config(config) + except Exception: + continue + self.skipTest(f"Cannot instantiate model with any Auto class for config {type(config).__name__}") + + def _get_tiny_config(self): + """Get config class and serialized dict for passing to spawned processes.""" + config = self.model_tester.get_config() + config.vocab_size = 256 + config.hidden_size = 64 + config.intermediate_size = 128 + if hasattr(config, "ffn_config"): + # Keep nested FFN projections consistent with resized hidden size. + if hasattr(config.ffn_config, "ffn_hidden_size"): + config.ffn_config.ffn_hidden_size = config.hidden_size + if hasattr(config.ffn_config, "hidden_size"): + config.ffn_config.hidden_size = config.intermediate_size + if hasattr(config, "num_attention_heads"): + config.num_attention_heads = 4 + if hasattr(config, "num_key_value_heads"): + config.num_key_value_heads = 4 + if hasattr(config, "vocab_size_per_layer_input"): + config.vocab_size_per_layer_input = config.vocab_size + # `to_diff_dict()` avoids nested config pollution (e.g. DBRX ffn_config receiving + # generic PretrainedConfig keys that its constructor rejects). + config_dict = config.to_diff_dict() + return type(config), config_dict + + def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs): + self._skip_if_fsdp_model_not_selected() + self._skip_if_insufficient_devices() + + config_class, config_dict = self._get_tiny_config() + func_args = (config_class, config_dict, *test_args) + + results_file = tempfile.mktemp(suffix=".json") + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + port = s.getsockname()[1] + + try: + mp.spawn( + _fsdp_global_wrapper, + args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file), + nprocs=self.fsdp_nproc_per_node, + ) + + with open(results_file) as f: + result = json.load(f) + finally: + if os.path.exists(results_file): + os.unlink(results_file) + + if result["error"] is not None: + self.fail(f"FSDP test '{test_name}' failed:\n{result['error']}") + + # ========================================================================= + # Test: get_transformer_block_classes (CPU, meta device) + # ========================================================================= + + @is_fsdp_test + def test_fsdp_plan_declared(self): + """The model exposes a non-empty `_fsdp_plan` derived from config + class-level overrides.""" + self._skip_if_fsdp_disabled() + self._skip_if_fsdp_model_not_selected() + start_time = time.perf_counter() + logger.info("[FSDP] Starting test: test_fsdp_plan_declared") + status = "FAIL" + try: + config = self.model_tester.get_config() + model = self._create_model_on_meta(config) + + plan = getattr(model, "_fsdp_plan", None) + self.assertTrue(plan, f"No _fsdp_plan declared for {type(model).__name__}") + status = "PASS" + finally: + logger.info("[FSDP] %s test: test_fsdp_plan_declared (%.1fs)", status, time.perf_counter() - start_time) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_sharding_structure_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_sharding_structure_untied", _test_fsdp2_sharding_structure_impl, False + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_sharding_structure_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_sharding_structure_tied", _test_fsdp2_sharding_structure_impl, True + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_save_load(self): + self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_auto_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_auto_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "auto" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_auto_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_auto_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "auto" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_manual_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_manual_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "manual" + ) + + @is_fsdp_test + @require_fsdp + def test_fsdp2_manual_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test( + "test_fsdp2_manual_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "manual" + ) diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 7720205c6fef..2e60939f9ead 100644 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -131,8 +131,8 @@ from torch import nn from transformers import MODEL_MAPPING + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan from transformers.integrations.accelerate import compute_module_sizes - from transformers.integrations.tensor_parallel import _get_parameter_tp_plan from transformers.modeling_utils import load_state_dict from transformers.pytorch_utils import id_tensor_storage diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 547bce7dacc4..d71c398c039e 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -16,8 +16,9 @@ from abc import ABC, abstractmethod from transformers import TorchAoConfig, set_seed -from transformers.distributed.configuration_utils import DistributedConfig -from transformers.integrations.tensor_parallel import _get_parameter_tp_plan +from transformers.distributed import DistributedConfig +from transformers.distributed.sharding_utils import _replicate_dtensor +from transformers.distributed.tensor_parallel import _get_parameter_tp_plan from transformers.testing_utils import ( is_tensor_parallel_test, is_torch_available, @@ -33,9 +34,26 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp + from torch.distributed.tensor import DTensor from torch.multiprocessing.spawn import ProcessRaisedException +def _to_local(tensor): + """Extract local tensor from DTensor, or return as-is for plain tensors.""" + if hasattr(tensor, "to_local"): + # NOTE(3outeille): With Sequence Parallelism, replicated params (e.g. norm weights) get Partial + # gradients — each rank holds only its contribution from its sequence shard. + # We must all-reduce (redistribute to Replicate) before extracting, otherwise + # we'd compare an incomplete gradient against the full reference. + # In the case of real training, we will always use SP + FSDP where the last will all-reduce the + # Partial gradients for us. + + if isinstance(tensor, DTensor) and any(not p.is_replicate() for p in tensor.placements): + tensor = _replicate_dtensor(tensor) + return tensor.to_local() + return tensor + + def _find_free_port(): """Find a free port by binding a socket and releasing it.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -112,50 +130,65 @@ def wrapper(*args, **kwargs): return _init_distributed_inner -def _load_tp_and_reference_models(model_path, model_class): +def _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=False): """Load TP model and non-TP reference model for comparison. Returns: tuple: (model_tp, model_ref, device) """ - model_tp = model_class.from_pretrained(model_path, tp_plan="auto") + tp_size = dist.get_world_size() + distributed_config = DistributedConfig(tp_size=tp_size, enable_sequence_parallel=enable_sequence_parallel) + model_tp = model_class.from_pretrained( + model_path, distributed_config=distributed_config, attn_implementation="sdpa" + ) dist.barrier() device = model_tp.device - model_ref = model_class.from_pretrained(model_path) + model_ref = model_class.from_pretrained(model_path, attn_implementation="sdpa") model_ref = model_ref.to(device) return model_tp, model_ref, device +def _get_active_tp_plan(model_tp): + distributed_config = getattr(model_tp.config, "distributed_config", None) + tp_plan = getattr(distributed_config, "tp_plan", None) + return tp_plan or getattr(model_tp, "_tp_plan", None) or {} + + def _verify_tp_sharding(rank, model_tp, model_ref): """Verify TP sharding by comparing parameter shapes between TP and reference models. + For DTensor params, uses the local tensor shape (not the global DTensor shape). + Returns: list: Names of sharded parameters """ world_size = dist.get_world_size() sharded_params = [] + tp_plan = _get_active_tp_plan(model_tp) for (name, param), (_, param_full) in zip(model_tp.named_parameters(), model_ref.named_parameters()): - if param.shape != param_full.shape: + # For DTensor params, get the local shape for comparison + param_local = _to_local(param) + if param_local.shape != param_full.shape: sharded_params.append(name) if rank == 0: - print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param.shape}") + print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param_local.shape}") # Verify sharding is correct - for dim in range(param.ndim): - if param.size(dim) != param_full.size(dim): - param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True) - if param_plan in ("packed_colwise",): + for dim in range(param_local.ndim): + if param_local.size(dim) != param_full.size(dim): + param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) + if param_plan == "packed_colwise": expected_size = param_full.size(dim) // world_size - assert param.size(dim) == expected_size, ( - f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param.size(dim)}" + assert param_local.size(dim) == expected_size, ( + f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param_local.size(dim)}" ) else: expected_size = (param_full.size(dim) + world_size - 1) // world_size - assert param.size(dim) <= expected_size, ( - f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param.size(dim)}" + assert param_local.size(dim) <= expected_size, ( + f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param_local.size(dim)}" ) break @@ -166,7 +199,7 @@ def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): """Implementation for comparing TP and non-TP model outputs.""" set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) _verify_tp_sharding(_rank, model_tp, model) @@ -179,7 +212,7 @@ def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): with torch.no_grad(): logits = model(input_ids).logits - logits_tp = model_tp(input_ids).logits + logits_tp = _to_local(model_tp(input_ids).logits) diff = (logits - logits_tp).abs() assert torch.allclose(logits, logits_tp, atol=atol, rtol=rtol), ( @@ -193,7 +226,8 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): """Implementation for comparing TP and non-TP model backward passes.""" set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) + model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) + tp_plan = _get_active_tp_plan(model_tp) model_tp.train() model.train() @@ -203,32 +237,35 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): set_seed(0) labels = torch.randint(0, vocab_size, (2, 64)).to(device) - loss = model(input_ids, labels=labels).loss + loss = model(input_ids, labels=labels, use_cache=False).loss loss.backward() - loss_tp = model_tp(input_ids, labels=labels).loss + loss_tp = model_tp(input_ids, labels=labels, use_cache=False).loss loss_tp.backward() - assert torch.allclose(loss, loss_tp, atol=atol, rtol=rtol), ( + loss_tp_local = _to_local(loss_tp) + assert torch.allclose(loss, loss_tp_local, atol=atol, rtol=rtol), ( f"TP and non-TP model losses differ. " - f"Non-TP loss: {loss.item()}, TP loss: {loss_tp.item()}, " - f"Diff: {(loss - loss_tp).abs().item()}" + f"Non-TP loss: {loss.item()}, TP loss: {loss_tp_local.item()}, " + f"Diff: {(loss - loss_tp_local).abs().item()}" ) # Compare gradients for matching parameters world_size = dist.get_world_size() + + # Debug: check tied weights and parameter alignment failed_grads = {} - for (name, param), (_, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): + for (name, param), (name_tp, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): if param.grad is not None and param_tp.grad is not None: grad = param.grad - grad_tp = param_tp.grad + grad_tp = _to_local(param_tp.grad) # Slice reference gradient to match local shard if parameter is sharded if grad.shape != grad_tp.shape: for dim in range(grad.ndim): if grad.size(dim) != grad_tp.size(dim): - param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True) - if param_plan in ("packed_colwise",): + param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) + if param_plan == "packed_colwise": # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: @@ -239,10 +276,15 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): break if not torch.allclose(grad.cpu(), grad_tp.cpu(), atol=atol, rtol=rtol): - failed_grads[name] = (grad.cpu() - grad_tp.cpu()).abs().max().item() + max_diff = (grad.cpu() - grad_tp.cpu()).abs().max().item() + ref_abs_max = grad.cpu().abs().max().item() + tp_abs_max = grad_tp.cpu().abs().max().item() + ratio = tp_abs_max / ref_abs_max if ref_abs_max > 0 else float("inf") + failed_grads[name] = (max_diff, ref_abs_max, tp_abs_max, ratio) assert not failed_grads, f"Gradients differ for {len(failed_grads)} parameter(s):\n" + "\n".join( - f" {name}: max diff = {diff}" for name, diff in failed_grads.items() + f" {name}: max_diff={v[0]:.6f}, ref_max={v[1]:.6f}, tp_max={v[2]:.6f}, tp/ref ratio={v[3]:.4f}" + for name, v in failed_grads.items() ) dist.barrier() @@ -274,7 +316,7 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new # Compare logits/scores at each generation step scores = torch.stack(output.scores) - scores_tp = torch.stack(output_tp.scores) + scores_tp = torch.stack([_to_local(s) for s in output_tp.scores]) diff = (scores - scores_tp).abs() assert torch.allclose(scores, scores_tp, atol=atol, rtol=rtol), ( @@ -283,9 +325,10 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new ) # Compare generated token sequences - assert torch.equal(output.sequences, output_tp.sequences), ( + sequences_tp = _to_local(output_tp.sequences) + assert torch.equal(output.sequences, sequences_tp), ( f"TP and non-TP model generated different token sequences (direct load path). " - f"Non-TP: {output.sequences.tolist()} | TP: {output_tp.sequences.tolist()}" + f"Non-TP: {output.sequences.tolist()} | TP: {sequences_tp.tolist()}" ) dist.barrier() @@ -297,7 +340,11 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t quantization_config = TorchAoConfig(Float8WeightOnlyConfig()) - model_tp = model_class.from_pretrained(model_path, tp_plan="auto", quantization_config=quantization_config) + model_tp = model_class.from_pretrained( + model_path, + distributed_config=DistributedConfig(tp_size=dist.get_world_size()), + quantization_config=quantization_config, + ) dist.barrier() device = model_tp.device @@ -347,9 +394,13 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t def _load_ep_and_reference_models(model_path, model_class): """Load EP model and non-EP reference model for comparison.""" + tp_size = dist.get_world_size() model_ep = model_class.from_pretrained( model_path, - distributed_config=DistributedConfig(enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=tp_size, + enable_expert_parallel=True, + ), ) dist.barrier() @@ -426,8 +477,8 @@ class TensorParallelTesterMixin(ABC): # Configuration (can be overridden per model) # ============================================================ tensor_parallel_size: int = 2 - tensor_parallel_atol: float = 1e-5 - tensor_parallel_rtol: float = 1e-5 + tensor_parallel_atol: float = 5e-3 + tensor_parallel_rtol: float = 5e-3 @property @abstractmethod @@ -562,18 +613,7 @@ def test_tp_generation_quantized(self): if not is_torchao_available(): self.skipTest("Test requires torchao") - config = self.model_tester.get_config() - model_class = self._get_tp_model_class() - max_new_tokens = 25 - - with tempfile.TemporaryDirectory() as tmp_dir: - set_seed(42) - model = model_class(config) - model.save_pretrained(tmp_dir, save_original_format=True) - - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_generation_quantized_impl)( - tmp_dir, model_class, max_new_tokens - ) + self.skipTest("Quantization is not currently supported with distributed training") @is_tensor_parallel_test def test_ep_forward(self): diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 9ac9810c9808..91e24988b81b 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -17,6 +17,7 @@ import torch import torch.nn as nn +from torch.distributed.tensor.placement_types import Shard from transformers import PreTrainedConfig, PreTrainedModel from transformers.conversion_mapping import ( @@ -39,11 +40,13 @@ convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, + spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig from transformers.utils.import_utils import is_triton_available from ..test_modeling_common import compare_state_dicts +from .test_distributed_sharding_utils import FakeMesh, _make_dtensor_shard_op class TestWeightGlobMatching(unittest.TestCase): @@ -225,6 +228,99 @@ def __init__(self, config, add_extra_moe=False, with_mlp=True): class TestConvertAndLoadStateDict(unittest.TestCase): + def test_dtensor_shard_aware_mixtral_conversion_uses_only_local_experts(self): + """Integration test: FSDP-sharded expert loading + WeightConverter. + + The problem: Mixtral has 8 experts. The checkpoint stores them separately:: + + experts.0.w1.weight (2x2) + experts.0.w3.weight (2x2) + experts.1.w1.weight (2x2) + experts.1.w3.weight (2x2) + + The model stores them packed into one tensor:: + + experts.gate_up_proj.weight (2, 4, 2) + ^ ^ ^ + | | +-- features + | +-- w1 (2) + w3 (2) concatenated + +-- num_experts + + The conversion (without FSDP) is: load all expert w1/w3 tensors, + MergeModulelist(dim=0) stacks experts, Concatenate(dim=1) joins w1+w3. + + With FSDP, Shard(0) splits the expert dim across ranks. Rank 0 owns + expert 0, rank 1 owns expert 1. So rank 0 should skip loading expert 1 + entirely -- not load it then discard it. + + What the test checks:: + + checkpoint files shard_tensor rank 0 gets + ---------------- ------------ ----------- + experts.0.w1 [[0,1],[2,3]] idx=0 -> kept [[0,1],[2,3]] + experts.1.w1 [[10,11],...] idx=1 -> None (not owned) + experts.0.w3 [[4,5],[6,7]] idx=0 -> kept [[4,5],[6,7]] + experts.1.w3 [[14,15],...] idx=1 -> None (not owned) + + WeightConverter then combines only the kept tensors:: + + MergeModulelist(dim=0): stack owned experts -> shape (1, 2, 2) each + Concatenate(dim=1): cat w1 + w3 along dim 1 + + gate_up_proj = [[[0,1],[2,3],[4,5],[6,7]]] shape (1, 4, 2) + ~~~~~~~~~~ ~~~~~~~~~~ + w1 w3 + + The key point: DtensorShardOperation.shard_tensor(tensor_idx=1) returns + None for rank 0, so the converter never even processes expert 1's data. + This saves memory during loading. + """ + shard_op = _make_dtensor_shard_op( + FakeMesh(shape=(2,), rank=0), + [Shard(0)], + param_shape=(2, 4, 2), + local_shape=(1, 4, 2), + ) + converter = WeightConverter( + ["experts.*.w1.weight", "experts.*.w3.weight"], + "experts.gate_up_proj.weight", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ) + + for idx, tensor in enumerate( + [ + torch.tensor([[0.0, 1.0], [2.0, 3.0]]), + torch.tensor([[10.0, 11.0], [12.0, 13.0]]), + ] + ): + converter.add_tensor( + "model.layers.0.experts.gate_up_proj.weight", + f"model.layers.0.experts.{idx}.w1.weight", + "experts.*.w1.weight", + spawn_materialize(None, tensor, device="cpu", dtype=None, sharding_op=shard_op, tensor_idx=idx), + ) + + for idx, tensor in enumerate( + [ + torch.tensor([[4.0, 5.0], [6.0, 7.0]]), + torch.tensor([[14.0, 15.0], [16.0, 17.0]]), + ] + ): + converter.add_tensor( + "model.layers.0.experts.gate_up_proj.weight", + f"model.layers.0.experts.{idx}.w3.weight", + "experts.*.w3.weight", + spawn_materialize(None, tensor, device="cpu", dtype=None, sharding_op=shard_op, tensor_idx=idx), + ) + + converted = converter.convert("model.layers.0.experts.gate_up_proj.weight") + + self.assertEqual(list(converted), ["model.layers.0.experts.gate_up_proj.weight"]) + torch.testing.assert_close( + converted["model.layers.0.experts.gate_up_proj.weight"], + torch.tensor([[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]]]), + ) + def test_moe_and_qkv_conversion(self): model = DummyRoot(PreTrainedConfig()) @@ -509,11 +605,11 @@ def __init__(self, config): model_state = model.state_dict() self.assertFalse(torch.allclose(raw_k, expected_k)) - torch.testing.assert_close(model_state["model.layers.0.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(model_state["model.layers.0.self_attn.v_proj.weight"], expected_v) + torch.testing.assert_close(model_state["layers.0.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(model_state["layers.0.self_attn.v_proj.weight"], expected_v) - q_weight_key = "model.layers.0.self_attn.q_proj.weight" - scale_key = "model.layers.0.self_attn.q_proj.weight_scale_inv" + q_weight_key = "layers.0.self_attn.q_proj.weight" + scale_key = "layers.0.self_attn.q_proj.weight_scale_inv" self.assertIn(scale_key, model_state) expected_dtype = torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else torch.int8 self.assertEqual(model_state[q_weight_key].dtype, expected_dtype) @@ -524,11 +620,14 @@ def __init__(self, config): torch.Size((out_dim // block_size[0], in_dim // block_size[1])), ) - dequant = Fp8Dequantize(block_size=block_size) + dequant = Fp8Dequantize(quantizer) dequantized_q = dequant.convert( - [model_state[q_weight_key], model_state[scale_key]], - context={"quantization_config": quantizer.quantization_config}, - ) + { + "weight$": [model_state[q_weight_key]], + "weight_scale_inv": [model_state[scale_key]], + }, + full_layer_name=q_weight_key, + )[q_weight_key] torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) def test_scoped_renaming_does_not_leak_to_sibling_or_parent(self): diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py new file mode 100644 index 000000000000..677f13ca2638 --- /dev/null +++ b/tests/utils/test_distributed_sharding_utils.py @@ -0,0 +1,462 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import shutil +import tempfile +import unittest + +import torch +from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard + +from transformers.distributed.sharding_utils import ( + DtensorShardOperation, + _find_strided_shard_placement_from_fused_params, +) + + +if torch.distributed.is_available(): + import torch.distributed as dist + import torch.distributed.checkpoint as dcp + import torch.multiprocessing as mp + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, distribute_tensor + + from transformers.distributed.sharding_utils import ( + fuse_optimizer_state, + get_fusion_metadata, + unfuse_optimizer_state, + ) + + +class FakeMesh: + """Fake multi-dimensional device mesh for testing DtensorShardOperation.""" + + def __init__(self, shape, rank, dim_names=None): + if isinstance(shape, int): + shape = (shape,) + self.shape = tuple(shape) + self.ndim = len(self.shape) + self.mesh_dim_names = dim_names or tuple(f"dim{i}" for i in range(self.ndim)) + # Compute nD coordinate (row-major: last dim changes fastest) + self._coord = [] + r = rank + for s in reversed(self.shape): + self._coord.insert(0, r % s) + r //= s + + def get_local_rank(self): + return self._coord[0] + + def get_coordinate(self): + return tuple(self._coord) + + def size(self): + result = 1 + for s in self.shape: + result *= s + return result + + def _is_current_rank_part_of_mesh(self): + return True + + def _sym_get_coordinate(self, dim): + return self._coord[dim] + + def __getitem__(self, name): + idx = self.mesh_dim_names.index(name) + return FakeMesh( + shape=(self.shape[idx],), + rank=self._coord[idx], + dim_names=(name,), + ) + + +def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): + """Build a DtensorShardOperation without requiring a real DTensor / distributed init. + + The axis-0 ownership cache is computed by mimicking + ``compute_local_shape_and_global_offset`` for the leading dim only: + locate the mesh dim that shards param dim 0 (if any) and use its local rank. + """ + op = object.__new__(DtensorShardOperation) + op.device_mesh = mesh + op.placements = tuple(placements) + op.param_ndim = len(param_shape) + op._axis0_offset = 0 + op._axis0_local_size = local_shape[0] + for mesh_dim, p in enumerate(placements): + if hasattr(p, "dim") and (p.dim % len(param_shape)) == 0: + sub = mesh[mesh.mesh_dim_names[mesh_dim]] if mesh.ndim > 1 else mesh + op._axis0_offset = sub.get_local_rank() * local_shape[0] + break + return op + + +class TestDtensorShardOperation(unittest.TestCase): + """Unit tests for DtensorShardOperation. + + See `DtensorShardOperation` in sharding_utils.py for the placement primer + and table of checkpoint layouts. The rest of this docstring covers the + test-specific conventions you need to write new cases here. + + Running example used throughout these tests: a stack of N MoE experts, + each of shape [in, out]. The full parameter shape is [N, in, out]. + Tests are parameterized so every rank in a (fake) mesh is checked. + + The checkpoint can store this param in two layouts, and shard_tensor + behaves differently for each: + + | Layout | tensor_idx | source.shape | + |------------------------------|-------------------|----------------| + | Single stacked tensor | None | [N, in, out] | + | N separate per-expert files | 0, 1, ..., N-1 | [in, out] | + + Single-tensor case: source has the full param shape, including axis 0. + shard_tensor returns this rank's slice along every sharded dim. + + Per-piece case: shard_tensor is called once per expert. Each call's + source is just that one expert's [in, out] tensor — note source has + one fewer dim than the param, because the axis-0 index lives in + `tensor_idx`, not in source.shape. shard_tensor returns: + - None, if this rank doesn't own expert `tensor_idx` along axis 0 + (the piece is then dropped by the caller, MergeModulelist / + Concatenate), + - the inner-dim slice otherwise. After all N calls, the caller + stacks the surviving slices along axis 0 to rebuild this rank's + local [n_local, in, out]. + """ + + def test_no_shard_placements_returns_full_copy(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor, # rank 0 — no shards, full copy + 1: tensor, # rank 1 — no shards, full copy + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_1D_shard(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[:2], # rank 0 — first half + 1: tensor[2:], # rank 1 — second half + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 4), local_shape=(2, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_1D_strided_shard(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[[0, 2]], # first piece of each group — rows {0, 2} + 1: tensor[[1, 3]], # second piece of each group — rows {1, 3} + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op( + mesh, [_StridedShard(dim=0, split_factor=2)], param_shape=(4, 4), local_shape=(2, 4) + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_shard_different_dims(self): + tensor = torch.arange(64).reshape(8, 8).float() + expected = { + 0: tensor[:4, :4], # top-left + 1: tensor[:4, 4:], # top-right + 2: tensor[4:, :4], # bottom-left + 3: tensor[4:, 4:], # bottom-right + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0), Shard(1)], param_shape=(8, 8), local_shape=(4, 4)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_shard_same_dim(self): + tensor = torch.arange(64).reshape(8, 8).float() + expected = { + 0: tensor[:2], # rows 0-1 + 1: tensor[2:4], # rows 2-3 + 2: tensor[4:6], # rows 4-5 + 3: tensor[6:8], # rows 6-7 + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0), Shard(0)], param_shape=(8, 8), local_shape=(2, 8)) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_strided_shard_same_dim(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: tensor[[0]], # row 0 + 1: tensor[[2]], # row 2 + 2: tensor[[1]], # row 1 + 3: tensor[[3]], # row 3 + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op( + mesh, + [_StridedShard(dim=0, split_factor=2), Shard(0)], + param_shape=(4, 4), + local_shape=(1, 4), + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_2D_strided_shard_different_dims(self): + tensor = torch.arange(16).reshape(4, 4).float() + expected = { + 0: torch.cat([tensor[:2, 0:1], tensor[:2, 2:3]], dim=1), # top rows, cols {0, 2} + 1: torch.cat([tensor[:2, 1:2], tensor[:2, 3:4]], dim=1), # top rows, cols {1, 3} + 2: torch.cat([tensor[2:, 0:1], tensor[2:, 2:3]], dim=1), # bottom rows, cols {0, 2} + 3: torch.cat([tensor[2:, 1:2], tensor[2:, 3:4]], dim=1), # bottom rows, cols {1, 3} + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op( + mesh, + [Shard(0), _StridedShard(dim=1, split_factor=2)], + param_shape=(4, 4), + local_shape=(2, 2), + ) + torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + + def test_moe_1D_shard_filters_by_axis0_ownership(self): + source = torch.ones(2, 2) + expected = { + 0: { + 0: source, # first owned + 1: source, # last owned + 2: None, # first not-owned (upper boundary, exclusive) + 3: None, # not owned + }, + 1: { + 0: None, # not owned + 1: None, # last not-owned (just below offset) + 2: source, # first owned (lower boundary, inclusive) + 3: source, # last owned + }, + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 2, 2), local_shape=(2, 2, 2)) + for tensor_idx, exp in expected[rank].items(): + with self.subTest(rank=rank, tensor_idx=tensor_idx): + shard = op.shard_tensor(source, tensor_idx=tensor_idx) + if exp is None: + self.assertIsNone(shard) + else: + torch.testing.assert_close(shard, exp) + + def test_moe_1D_strided_shard_on_inner_dim_degrades_to_contiguous(self): + source = torch.arange(8).reshape(4, 2).float() + expected = { + 0: source[:2], # rank 0 — first half (strided silently degraded to contiguous) + 1: source[2:], # rank 1 — second half + } + for rank in range(2): + mesh = FakeMesh(shape=(2,), rank=rank) + op = _make_dtensor_shard_op( + mesh, + [_StridedShard(dim=1, split_factor=2)], + param_shape=(8, 8, 2), + local_shape=(8, 4, 2), + ) + torch.testing.assert_close(op.shard_tensor(source, tensor_idx=0), expected[rank], msg=f"rank {rank}") + + def test_moe_2D_shard_on_axis0_and_inner_dim_slices_inner(self): + source = torch.arange(8).reshape(4, 2).float() + expected = { + 0: source[:2], # owned; inner Shard(1) → source dim 0 first half + 1: source[2:], # owned; inner Shard(1) → source dim 0 second half + 2: None, # not owned (axis-0 ownership filter) + 3: None, # not owned + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(0), Shard(1)], param_shape=(4, 4, 2), local_shape=(2, 2, 2)) + shard = op.shard_tensor(source, tensor_idx=1) + if expected[rank] is None: + self.assertIsNone(shard, msg=f"rank {rank}") + else: + torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + + def test_moe_2D_shard_with_negative_dim_indices(self): + source = torch.arange(8).reshape(4, 2).float() + expected = { + 0: source[:, :1], # owned; inner Shard(-1) → source dim 1, first half + 1: source[:, 1:], # owned; inner Shard(-1) → source dim 1, second half + 2: None, # not owned + 3: None, # not owned + } + for rank in range(4): + mesh = FakeMesh(shape=(2, 2), rank=rank) + op = _make_dtensor_shard_op(mesh, [Shard(-3), Shard(-1)], param_shape=(4, 4, 2), local_shape=(2, 4, 1)) + shard = op.shard_tensor(source, tensor_idx=1) + if expected[rank] is None: + self.assertIsNone(shard, msg=f"rank {rank}") + else: + torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + + def test_strided_intervals(self): + # Direct tests for _strided_intervals(intervals, rank, world_size, split_factor). + # Keys: (input_interval, rank, world_size, split_factor) -> expected output list. + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(8,), local_shape=(4,)) + expected = { + # Even (0, 8) sf=2 ws=2 -> groups (0,4) (4,8); each rank takes half of each + ((0, 8), 0, 2, 2): [(0, 2), (4, 6)], + ((0, 8), 1, 2, 2): [(2, 4), (6, 8)], + # Uneven (0, 7) sf=2 -> group 0 = (0,4), group 1 = (4,7) -> size 3 + ((0, 7), 0, 2, 2): [(0, 2), (4, 6)], # rank 0 -> half of each group + ((0, 7), 1, 2, 2): [(2, 4), (6, 7)], # rank 1's piece in group 1 is 1 elem wide + # split_factor=1 collapses to a single group -> contiguous behavior + ((0, 4), 0, 2, 1): [(0, 2)], + # split_factor=4 on size-2: groups (0,1), (1,2), (2,2), (3,2) -> last 2 empty -> skipped + ((0, 2), 0, 2, 4): [(0, 1), (1, 2)], + } + for (interval, rank, ws, sf), exp in expected.items(): + with self.subTest(interval=interval, rank=rank, ws=ws, sf=sf): + self.assertEqual(op._strided_intervals([interval], rank=rank, world_size=ws, split_factor=sf), exp) + + def test_contiguous_intervals(self): + # Direct tests for _contiguous_intervals(intervals, rank, world_size). + # Keys: (input_intervals, rank, world_size) -> expected output list. + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(8,), local_shape=(4,)) + expected = { + # Single interval, even split -> rank 0 -> first 4 elems, rank 1 -> last 4 elems + (((0, 8),), 0, 2): [(0, 4)], + (((0, 8),), 1, 2): [(4, 8)], + # Uneven: size 5 / 2 ranks -> rank 0 -> first 3 elems, rank 1 -> last 2 elems + (((0, 5),), 0, 2): [(0, 3)], + (((0, 5),), 1, 2): [(3, 5)], + # Empty: size 3 / 4 ranks -> rank 3 -> nothing + (((0, 3),), 3, 4): [], + # Multi-input intervals (8 elems): rank -> takes its slice from whichever interval(s) cover it + (((0, 4), (10, 14)), 0, 2): [(0, 4)], # rank 0 -> first 4 elems = first interval entirely + (((0, 4), (10, 14)), 1, 2): [(10, 14)], # rank 1 -> last 4 elems = second interval entirely + (((0, 4), (10, 14)), 2, 4): [(10, 12)], # ws=4, rank 2 -> cuts mid-input-interval + } + for (intervals, rank, ws), exp in expected.items(): + with self.subTest(intervals=intervals, rank=rank, ws=ws): + self.assertEqual(op._contiguous_intervals(list(intervals), rank=rank, world_size=ws), exp) + + def test_slice_and_cat(self): + # Direct tests for _slice_and_cat(source, intervals, device, dtype). + tensor = torch.arange(64).reshape(8, 8).float() + mesh = FakeMesh(shape=(2,), rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(8, 8), local_shape=(4, 8)) + expected = { + # Fast path: every dim is single-interval -> one slice read, no cat + "fast_path": ([[(0, 4)], [(0, 8)]], tensor[:4, :]), + # Cat along dim 1: two disjoint col ranges -> read separately and concat + "cat_dim1": ([[(0, 4)], [(0, 2), (4, 6)]], torch.cat([tensor[:4, :2], tensor[:4, 4:6]], dim=1)), + # Cat along dim 0: two disjoint row ranges -> read separately and concat + "cat_dim0": ([[(0, 2), (4, 6)], [(0, 8)]], torch.cat([tensor[:2, :], tensor[4:6, :]], dim=0)), + } + for case, (intervals, exp) in expected.items(): + with self.subTest(case=case): + torch.testing.assert_close(op._slice_and_cat(tensor, intervals, None, None), exp) + + # Reject: two dims with disjoint ranges -> would require an outer-product of reads. + with self.assertRaises(ValueError): + op._slice_and_cat(tensor, [[(0, 2), (4, 6)], [(0, 2), (4, 6)]], None, None) + + result = op._slice_and_cat(tensor, [[(0, 4)], [(0, 8)]], None, torch.float16) + self.assertEqual(result.dtype, torch.float16) + + +class TestFindStridedShardPlacementFromFusedParams(unittest.TestCase): + def test_find_strided_shard_placement_from_fused_params(self): + expected = { + # Plain placements — no _StridedShard, nothing to do + (Replicate(),): None, + (Shard(0),): None, + (Shard(0), Shard(2)): None, + # Uncomposed _StridedShard — TP-only fused gate||up (DCP can't encode) + (_StridedShard(0, split_factor=2),): _StridedShard(0, split_factor=2), + (_StridedShard(1, split_factor=4),): _StridedShard(1, split_factor=4), + # _StridedShard on a different tensor dim than the other Shard — still uncomposed + (Shard(0), _StridedShard(1, split_factor=2)): _StridedShard(1, split_factor=2), + (_StridedShard(2, split_factor=2), Shard(0)): _StridedShard(2, split_factor=2), + # _StridedShard composed with another Shard on the SAME tensor dim — DCP-friendly + (_StridedShard(0, split_factor=2), Shard(0)): None, + (Shard(0), _StridedShard(0, split_factor=2)): None, + } + for placements, exp in expected.items(): + with self.subTest(placements=placements): + self.assertEqual(_find_strided_shard_placement_from_fused_params(placements), exp) + + +def _optimizer_state_checkpointing_e2e_worker(rank, world_size, port, ckpt_dir): + # 1. Init a 4-rank CPU process group + 2x2 (fsdp, tp) mesh. + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("fsdp", "tp")) + + # 2. Build a fused DTensor with the Mixtral gate_up_proj placement: + # shape (num_experts=4, 2·intermediate=16, hidden=8), + # placements (Shard(0), _StridedShard(1, sf=2)). + full = torch.arange(4 * 16 * 8, dtype=torch.float32).reshape(4, 16, 8) + dt = distribute_tensor(full, mesh, [Shard(0), Shard(1)]) + dt = DTensor.from_local( + dt._local_tensor.clone(), + mesh, + (Shard(0), _StridedShard(1, split_factor=2)), + run_check=False, + ) + + # 3. Wrap it the way `get_optimizer_state_dict` would. + fqn = "model.layers.0.mlp.experts.gate_up_proj" + osd = { + "state": {fqn: {"exp_avg": dt, "step": torch.tensor(7.0)}}, + "param_groups": [{"lr": 1e-4}], + } + + # 4. Snapshot the rank-local buffer for later bit-exact comparison. + before = dt._local_tensor.clone() + + # 5. unfuse → DCP save → DCP load → fuse. + fusion_metadata = get_fusion_metadata(osd) + assert set(fusion_metadata) == {fqn}, f"expected one fused param, got {set(fusion_metadata)}" + unfuse_optimizer_state(osd, fusion_metadata) + dcp.save({"optimizer": osd}, checkpoint_id=ckpt_dir) + dist.barrier() + dcp.load({"optimizer": osd}, checkpoint_id=ckpt_dir) + fuse_optimizer_state(osd, fusion_metadata) + + # 6. Verify the placement is restored and the rank-local data is bit-exact. + after = osd["state"][fqn]["exp_avg"] + assert tuple(after.placements) == (Shard(0), _StridedShard(1, split_factor=2)), after.placements + assert torch.equal(after._local_tensor, before), f"rank {rank}: local data drifted after round-trip" + + dist.destroy_process_group() + + +@unittest.skipUnless(torch.distributed.is_available(), "Requires torch.distributed (gloo backend).") +class TestOptimizerStateCheckpointing(unittest.TestCase): + def test_optimizer_state_checkpointing_e2e(self): + tmp = tempfile.mkdtemp(prefix="hf_optimizer_state_checkpointing_") + try: + mp.spawn(_optimizer_state_checkpointing_e2e_worker, args=(4, 29500, tmp), nprocs=4, join=True) + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 025bf3ca79ff..5c26f44f7736 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -59,6 +59,7 @@ is_torch_available, logging, ) +from transformers.distributed import DistributedConfig from transformers.modeling_flash_attention_utils import is_flash_attn_available from transformers.models.mistral.modeling_mistral import MistralModel from transformers.testing_utils import ( @@ -432,6 +433,43 @@ def test_get_total_byte_count_does_not_require_process_group(self): self.assertIn(torch.device("cpu"), total_byte_count) self.assertGreater(total_byte_count[torch.device("cpu")], 0) + def test_model_from_pretrained_fsdp_distributes_before_loading(self): + model = GPT2LMHeadModel(GPT2Config(n_layer=1, n_head=2, n_embd=8, n_positions=8, n_ctx=8, vocab_size=32)) + + with tempfile.TemporaryDirectory() as tmp_dir: + model.save_pretrained(tmp_dir) + call_order = [] + + fake_mesh = mock.Mock() + fake_mesh.mesh_dim_names = ("fsdp",) + fake_mesh.ndim = 1 + fake_mesh.device_type = "cpu" + + def fake_apply_fsdp(model, fsdp_mesh, fsdp_plan): + call_order.append("distribute") + self.assertIs(fsdp_mesh, fake_mesh) + self.assertIsNone(fsdp_plan) + return model + + def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None): + call_order.append("load") + self.assertIs(load_config.device_mesh, fake_mesh) + return mock.Mock(), None + + with ( + patch("transformers.modeling_utils.init_device_mesh", return_value=fake_mesh), + patch("transformers.distributed.utils.apply_fully_shard_data_parallel", side_effect=fake_apply_fsdp), + patch.object(GPT2LMHeadModel, "_load_pretrained_model", side_effect=fake_load_pretrained_model), + patch.object( + GPT2LMHeadModel, + "_finalize_model_loading", + side_effect=lambda model, load_config, loading_info: loading_info, + ), + ): + GPT2LMHeadModel.from_pretrained(tmp_dir, distributed_config=DistributedConfig(fsdp_size=2)) + + self.assertEqual(call_order, ["distribute", "load"]) + def test_hub_retry(self): @hub_retry(max_attempts=2) def test_func(): diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index a138ef2eaacb..f09480a59875 100644 --- a/utils/tests_fetcher.py +++ b/utils/tests_fetcher.py @@ -1099,6 +1099,7 @@ def parse_commit_message(commit_message: str) -> dict[str, bool]: "tests_non_model": r"tests/[^/]*?/test_.*\.py", "tests_training_ci": r"tests/models/.*/test_modeling_.*", "tests_tensor_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/tensor_parallel(?:/test_tensor_parallel\.py)?)", + "tests_fsdp_ci": r"(tests/models/.*/test_modeling_.*|tests/test_fsdp_mixin\.py)", } From 863e12ba3c17a25180849e992d87a9f813746f0c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 27 May 2026 18:28:25 +0000 Subject: [PATCH 2/9] not unifying but still better than before --- .../distributed/tensor_parallel.py | 49 ++++++++++++------- 1 file changed, 32 insertions(+), 17 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 126b404594f2..458da9787884 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -116,21 +116,33 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") -class TensorParallelStyle(ParallelStyle): - """Base class for transformers TP styles. Installs the pre / around / post - forward hooks. Subclasses that need to shard params override `_apply` to - wrap them as DTensor placeholders before calling `super()._apply(...)`. +class CustomParallelStyle(ParallelStyle): + """Base class for custom TP styles (MoE, packed linear, etc.). - Param wrapping runs on meta (the model is on meta when `apply_tensor_parallel` - is invoked); `distribute_tensor` on meta builds metadata only — no collective. + Apply runs in two phases (subclasses should not override _apply): + + 1. shard_parameters(module, mesh) — meta-time DTensor placeholders (optional) + 2. _install_forward(module, mesh) — pre / around / post forward transforms + + Unlike PyTorch registry entries (ColwiseParallel, SequenceParallel, …) which use + distribute_module or forward hooks via their own _apply, this base wraps + module.forward to support *args, **kwargs and an optional context manager. + + Param wrapping runs on meta (the model is on meta when apply_tensor_parallel + is invoked); distribute_tensor on meta builds metadata only — no collective. Real data flows in later, async, via DtensorShardOperation during load. - Forward-time hooks (override what you need): + Override what you need: + - shard_parameters(module, mesh) — wrap params as DTensor placeholders - transform_inputs_pre_forward(module, args, kwargs, mesh) → (args, kwargs) - context_around_forward(module) → context manager wrapping the call - transform_output_post_forward(module, output, mesh) → output """ + def shard_parameters(self, module, mesh): + """Wrap selected params as DTensor placeholders. Default: no-op.""" + pass + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): return args, kwargs @@ -140,7 +152,8 @@ def context_around_forward(self, module): def transform_output_post_forward(self, module, output, mesh): return output - def _apply(self, module, mesh): + def _install_forward(self, module, mesh): + """Install pre / around / post transforms by replacing module.forward.""" original_forward = module.forward def tp_forward(*args, **kwargs): @@ -152,8 +165,12 @@ def tp_forward(*args, **kwargs): module.forward = tp_forward return module + def _apply(self, module, mesh): + self.shard_parameters(module, mesh) + return self._install_forward(module, mesh) + -class PrepareModuleInputOutput(TensorParallelStyle): +class PrepareModuleInputOutput(CustomParallelStyle): """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). Used for MoE blocks with SP: the input sequence is gathered before routing, @@ -244,7 +261,7 @@ def _swap_dtensor_params_for_local(module): module.register_parameter(name, param) -class PackedColwiseParallel(TensorParallelStyle): +class PackedColwiseParallel(CustomParallelStyle): """Column-wise parallel style for fused linear weights packed along the output dimension.""" def __init__( @@ -278,7 +295,7 @@ def transform_output_post_forward(self, module, output, mesh): output, mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False ) - def _apply(self, module, mesh): + def shard_parameters(self, module, mesh): if not isinstance(module, torch.nn.Linear): raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") # Wrap weight + bias as DTensor placeholders. Runs on meta — @@ -292,7 +309,6 @@ def _apply(self, module, mesh): distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, ) - return super()._apply(module, mesh) def __repr__(self) -> str: return ( @@ -323,14 +339,14 @@ def backward(ctx, grad): return grad, None -class MoEExpertsParallel(TensorParallelStyle): +class MoEExpertsParallel(CustomParallelStyle): """Tensor-parallel style for MoE expert modules. - Shards expert weights as DTensors, then wraps the module's ``forward`` so + Shards expert weights as DTensors, then wraps the module's forward so that grouped_mm (which needs plain tensors) works transparently. Lifecycle phases: - 1. _apply — wrap each expert weight named in shard_plan as a DTensor + 1. shard_parameters — wrap each expert weight named in shard_plan as a DTensor placeholder with the declared placement. 2. transform_inputs_pre_forward — localize hidden_states (Replicate→local, gives us an all-reduce on the backward gradient for free), then fix @@ -349,7 +365,7 @@ def __init__(self, output_layouts=None, shard_plan=None): self.output_layouts = output_layouts or Replicate() self._moe_shard_plan = shard_plan or {} - def _apply(self, module, mesh): + def shard_parameters(self, module, mesh): # Wrap each expert weight as a DTensor placeholder. Runs on meta — # distribute_tensor builds metadata only, no collective. for name, placement in self._moe_shard_plan.items(): @@ -360,7 +376,6 @@ def _apply(self, module, mesh): distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, ) - return super()._apply(module, mesh) def transform_inputs_pre_forward(self, module, args, kwargs, mesh): hidden_states, top_k_index, top_k_weights = args From f914e301d69cfb25adc634fae9adb470955956e8 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 28 May 2026 18:08:58 +0000 Subject: [PATCH 3/9] mellum now inherits tp and sp plan from qwen3_moe --- .../models/mellum/configuration_mellum.py | 36 ++++++++++++++----- .../models/mellum/modeling_mellum.py | 4 ++- 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/src/transformers/models/mellum/configuration_mellum.py b/src/transformers/models/mellum/configuration_mellum.py index 49deff0e2e69..15d24bbc4d6d 100644 --- a/src/transformers/models/mellum/configuration_mellum.py +++ b/src/transformers/models/mellum/configuration_mellum.py @@ -54,15 +54,29 @@ class MellumConfig(PreTrainedConfig): "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", } # Expert-only EP plan: only shards MoE experts, not attention. # Attention is left unsharded — FSDP2 handles attention weight distribution. @@ -71,7 +85,7 @@ class MellumConfig(PreTrainedConfig): "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -79,6 +93,12 @@ class MellumConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } + vocab_size: int = 98304 hidden_size: int = 2304 intermediate_size: int = 7168 diff --git a/src/transformers/models/mellum/modeling_mellum.py b/src/transformers/models/mellum/modeling_mellum.py index bc206be988ec..1b4ba4f1d725 100644 --- a/src/transformers/models/mellum/modeling_mellum.py +++ b/src/transformers/models/mellum/modeling_mellum.py @@ -633,8 +633,10 @@ def load_balancing_loss_func( @auto_docstring class MellumForCausalLM(MellumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): super().__init__(config) From 2d01eb3f974c30e514b3a2f1de429750232d24d9 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Thu, 28 May 2026 18:50:00 +0000 Subject: [PATCH 4/9] handle dense + sparse mixing for mellum model in SP plan --- .../models/mellum/configuration_mellum.py | 30 ++++++++++++--- .../models/mellum/modular_mellum.py | 38 +++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/src/transformers/models/mellum/configuration_mellum.py b/src/transformers/models/mellum/configuration_mellum.py index 15d24bbc4d6d..f29681a631f3 100644 --- a/src/transformers/models/mellum/configuration_mellum.py +++ b/src/transformers/models/mellum/configuration_mellum.py @@ -60,6 +60,8 @@ class MellumConfig(PreTrainedConfig): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", } + + # Per-layer MLP entries are filled in by `_update_sp_plan` (dense vs sparse). base_model_sp_plan = { "embed_tokens": "vocab_reduce_scatter", "layers.*.input_layernorm": "activation", @@ -71,11 +73,6 @@ class MellumConfig(PreTrainedConfig): "layers.*.self_attn.q_norm": "activation_seq_dim_2", "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.post_attention_layernorm": "activation", - "layers.*.mlp": "module_allgather_split", - "layers.*.mlp.experts": "moe_experts_allreduce", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise_reduce_scatter", "norm": "activation", } # Expert-only EP plan: only shards MoE experts, not attention. @@ -140,11 +137,34 @@ def __post_init__(self, **kwargs): "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, } + self._update_sp_plan() + super().__post_init__( **kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}, ) + def _update_sp_plan(self): + """Set per-layer SP entries depending on whether the MLP is dense or sparse.""" + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + def convert_rope_params_to_dict(self, **kwargs): # No need to handle BC for new models, because they have no old-format `rope_scaling` return kwargs diff --git a/src/transformers/models/mellum/modular_mellum.py b/src/transformers/models/mellum/modular_mellum.py index 889b4790088c..2fb71f4ebfa8 100644 --- a/src/transformers/models/mellum/modular_mellum.py +++ b/src/transformers/models/mellum/modular_mellum.py @@ -54,6 +54,21 @@ class MellumConfig(Qwen3MoeConfig): model_type = "mellum" + # Per-layer MLP entries are filled in by `_update_sp_plan` (dense vs sparse). + base_model_sp_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "norm": "activation", + } + vocab_size: int = 98304 hidden_size: int = 2304 intermediate_size: int = 7168 @@ -84,12 +99,35 @@ def __post_init__(self, **kwargs): "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, } + self._update_sp_plan() + PreTrainedConfig.__post_init__( self, **kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}, ) + def _update_sp_plan(self): + """Set per-layer SP entries depending on whether the MLP is dense or sparse.""" + self.base_model_sp_plan = self.base_model_sp_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + def convert_rope_params_to_dict(self, **kwargs): # No need to handle BC for new models, because they have no old-format `rope_scaling` return kwargs From fb95e2644ed0d0d5ee0d9245171a9ed3486971dc Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 16 Jun 2026 03:34:23 +0000 Subject: [PATCH 5/9] check-repository-consistency --- .../configuration_deepseek_ocr2.py | 1 - .../deepseek_ocr2/modular_deepseek_ocr2.py | 14 +++++++++++ .../deepseek_v32/modular_deepseek_v32.py | 5 +++- .../modular_diffusion_gemma.py | 25 +++++++++++++++++++ .../gemma4_unified/modular_gemma4_unified.py | 5 ++++ .../modular_gemma4_unified_assistant.py | 2 ++ .../glm_moe_dsa/configuration_glm_moe_dsa.py | 12 ++++----- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 2 +- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 12 ++++++--- .../minimax_m3_vl/modular_minimax_m3_vl.py | 22 ++++++++++++++++ .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 1 + 11 files changed, 89 insertions(+), 12 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py index 2d9c601c1f2a..34e666e7a260 100644 --- a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py @@ -88,7 +88,6 @@ class DeepseekOcr2VisionEncoderConfig(PreTrainedConfig): model_type = "deepseek_ocr2_encoder" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `DeepseekOcr2VisionEncoder` base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index cddbd6f2151e..a9ebd9471616 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -533,6 +533,18 @@ class DeepseekOcr2VisionEncoderConfig(Qwen2Config): ```""" base_config_key = "encoder_config" + # Default tensor parallel plan for base model `DeepseekOcr2VisionEncoder` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_sp_plan = AttributeError() + base_model_fsdp_plan = AttributeError() @auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") @@ -600,6 +612,8 @@ class DeepseekOcr2TextConfig(DeepseekV2Config): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } + base_model_sp_plan = AttributeError() + base_model_fsdp_plan = AttributeError() # Remove unused attributes inherited from DeepseekV2Config first_k_dense_replace = AttributeError() diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py index 77be4aa9c943..baf210824e85 100644 --- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py @@ -98,6 +98,7 @@ class DeepseekV32Config(Glm4MoeLiteConfig, RotaryEmbeddingConfigMixin): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } + base_model_fsdp_plan = AttributeError() attribute_map = {"num_local_experts": "num_experts"} @@ -369,7 +370,9 @@ class DeepseekV32Model(DeepseekV3Model): class DeepseekV32ForCausalLM(DeepseekV3ForCausalLM): - pass + _tp_plan = {"lm_head": "colwise_gather_output"} + _sp_plan = AttributeError() + _fsdp_plan = AttributeError() __all__ = [ diff --git a/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py index 6f5759cf3d23..07d9d7e2e3bf 100644 --- a/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py +++ b/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py @@ -89,11 +89,36 @@ class DiffusionGemmaTextConfig(Gemma4TextConfig): model_type = "diffusion_gemma_text" final_logit_softcapping = 30.0 + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + "layers.*.experts.gate_up_proj": "packed_colwise", + "layers.*.experts.down_proj": "rowwise", + "layers.*.experts": "moe_tp_experts", + } + base_model_ep_plan = { + # EP plan for google/gemma-4-26B-A4B-it: do not tp in attention (num_global_key_value_heads=2 too small to partition) + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + "layers.*.router": "ep_router", + "layers.*.experts.gate_up_proj": "grouped_gemm", + "layers.*.experts.down_proj": "grouped_gemm", + "layers.*.experts": "moe_tp_experts", + } base_model_pp_plan = { "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + base_model_fsdp_plan = AttributeError() enable_moe_block = AttributeError() attention_k_eq_v = AttributeError() diff --git a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py index 17d2cd107c20..1967c45b839b 100644 --- a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py +++ b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py @@ -456,6 +456,7 @@ class Gemma4UnifiedTextConfig(Gemma4TextConfig): "layers.*.mlp.down_proj": "rowwise", } base_model_ep_plan = None # No MoE + base_model_fsdp_plan = AttributeError() sliding_window: int = 1024 use_bidirectional_attention: Literal["all", "vision"] | None = "vision" @@ -757,6 +758,10 @@ def forward( class Gemma4UnifiedForCausalLM(Gemma4ForCausalLM): + _tp_plan = {"lm_head": "colwise_gather_output"} + _sp_plan = AttributeError() + _fsdp_plan = AttributeError() + def forward( self, input_ids: torch.LongTensor | None = None, diff --git a/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py b/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py index 9b9e619c2a53..59f4f3c74e73 100644 --- a/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py +++ b/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py @@ -16,6 +16,8 @@ class Gemma4UnifiedAssistantForCausalLM(Gemma4AssistantForCausalLM): + _fsdp_plan = AttributeError() + def forward(**super_kwargs): r""" shared_kv_states (`dict[str, torch.Tensor` of shape `(batch_size, 1, q_len, kv_len)`, *optional*): diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index c0758cbba377..0cac376ea2fb 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -79,12 +79,6 @@ class GlmMoeDsaConfig(PreTrainedConfig, RotaryEmbeddingConfigMixin): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } - - base_model_fsdp_plan = { - "embed_tokens": "free_full_weight", - "layers.*": "free_full_weight", - "norm": "keep_full_weight", - } attribute_map = { "num_local_experts": "n_routed_experts", } @@ -129,6 +123,12 @@ class GlmMoeDsaConfig(PreTrainedConfig, RotaryEmbeddingConfigMixin): head_dim: int = 64 first_k_dense_replace: int = 3 layer_types: list[str] | None = None + + base_model_fsdp_plan = { + "embed_tokens": "free_full_weight", + "layers.*": "free_full_weight", + "norm": "keep_full_weight", + } # `"full"` runs the indexer, `"shared"` reuses the previous full layer's index mask. indexer_types: list[str] | None = None diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index e120430ec392..dbf6fc56b55a 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -782,8 +782,8 @@ def forward( class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} - _sp_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _sp_plan = {"lm_head": "colwise_loss_parallel"} _fsdp_plan = {"lm_head": "keep_full_weight"} def __init__(self, config): diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index d76ddcb688a5..46845af56d5f 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -90,13 +90,17 @@ class GlmMoeDsaConfig(DeepseekV32Config): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", - "norm": "keep_full_weight" + "norm": "keep_full_weight", } - attribute_map = { "num_local_experts": "n_routed_experts", } @@ -352,7 +356,9 @@ def forward( class GlmMoeDsaForCausalLM(DeepseekV32ForCausalLM): - pass + _tp_plan = {"lm_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _fsdp_plan = {"lm_head": "keep_full_weight"} __all__ = [ diff --git a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py index 173795deb2be..fed287364099 100644 --- a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py @@ -93,6 +93,17 @@ class MiniMaxM3VLTextConfig(MiniMaxM2Config): model_type = "minimax_m3_vl_text" base_config_key = "text_config" + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise_gather_output", + "layers.*.self_attn.k_proj": "colwise_gather_output", + "layers.*.self_attn.v_proj": "colwise_gather_output", + "layers.*.self_attn.o_proj": "rowwise_split_input", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + base_model_sp_plan = AttributeError() + base_model_fsdp_plan = AttributeError() base_model_ep_plan = { "layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", @@ -101,6 +112,14 @@ class MiniMaxM3VLTextConfig(MiniMaxM2Config): "layers.*.mlp.experts.down_proj_scale_inv": "grouped_gemm", "layers.*.mlp.experts": "moe_tp_experts", } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + attribute_map = { + "num_experts": "num_local_experts", + } hidden_size: int = 6144 intermediate_size: int = 3072 @@ -711,6 +730,9 @@ def forward( class MiniMaxM3VLForCausalLM(MiniMaxM2ForCausalLM): config: MiniMaxM3VLTextConfig _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _sp_plan = AttributeError() + _fsdp_plan = AttributeError() def __init__(self, config: MiniMaxM3VLTextConfig): super().__init__(config) diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py index 424909e88a35..edbac33dc6ec 100644 --- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py @@ -245,6 +245,7 @@ def __init__(self, config): class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): + config: Qwen3_5MoeConfig _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} From 61d3cbd1cb4494d35242d6bbff8abd5bb5745279 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 16 Jun 2026 04:00:38 +0000 Subject: [PATCH 6/9] check-code-quality --- src/transformers/cli/serving/utils.py | 8 ++++---- src/transformers/generation/utils.py | 17 +++++++++++++++-- src/transformers/modeling_utils.py | 12 +++++++++--- 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/src/transformers/cli/serving/utils.py b/src/transformers/cli/serving/utils.py index 8aedfb45c0b8..0b2c909313ec 100644 --- a/src/transformers/cli/serving/utils.py +++ b/src/transformers/cli/serving/utils.py @@ -491,7 +491,7 @@ def __init__( self._tokenizer = tokenizer self._loop = loop self._queue = queue - self._decode_stream = DecodeStream([], skip_special_tokens) + self._decode_stream: DecodeStream = DecodeStream([], skip_special_tokens) self._stc_id = tool_config["stc_id"] if tool_config else None self._etc_id = tool_config["etc_id"] if tool_config else None self._inside_tool_call = False @@ -523,7 +523,7 @@ def put(self, value: "torch.Tensor") -> None: is_start_or_end_token = _advance_thinking_state(self, token_id) - text = self._decode_stream.step(self._tokenizer, token_id) + text = self._decode_stream.step(self._tokenizer, token_id) # ty:ignore[unresolved-attribute] if text is None or self._inside_tool_call or token_id == self._etc_id or is_start_or_end_token: continue if self._inside_thinking: @@ -575,7 +575,7 @@ def __init__( self._loop = loop self._queue = queue self._tokenizer = tokenizer - self._decode_stream = DecodeStream([], True) + self._decode_stream: DecodeStream = DecodeStream([], True) self._stc_id = tool_config["stc_id"] if tool_config else None self._etc_id = tool_config["etc_id"] if tool_config else None self._inside_tool_call = False @@ -602,7 +602,7 @@ def put(self, output: "GenerationOutput") -> None: is_start_or_end_token = _advance_thinking_state(self, token_id) - text = self._decode_stream.step(self._tokenizer, token_id) + text = self._decode_stream.step(self._tokenizer, token_id) # ty:ignore[unresolved-attribute] if text is None or self._inside_tool_call or token_id == self._etc_id or is_start_or_end_token: continue if self._inside_thinking: diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 88cdafa71d35..ed14b60fbfa8 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -118,6 +118,18 @@ logger = logging.get_logger(__name__) + +def _get_dist_world_size() -> int: + if ( + dist.is_available() + and hasattr(dist, "is_initialized") + and dist.is_initialized() + and hasattr(dist, "get_world_size") + ): + return dist.get_world_size() + return 1 + + if is_accelerate_available(): from accelerate.hooks import AlignDevicesHook, add_hook_to_module @@ -2165,7 +2177,7 @@ def _extract_generation_mode_kwargs( "assistant_model": assistant_model, "streamer": streamer, } - world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 + world_size = _get_dist_world_size() generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2612,7 +2624,8 @@ def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, # The following logic allows an early break if all peers finished generating their sequence this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) # send 0.0 if we finished, 1.0 otherwise - dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) + if hasattr(dist, "all_reduce") and hasattr(dist, "ReduceOp"): + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # did all peers finish? the reduced sum will be 0.0 then if this_peer_finished_flag.item() == 0.0: return False diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 5f7105317a32..2a2bd200e12c 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -209,6 +209,12 @@ def _get_torch_distributed_world_size() -> int: return torch.distributed.get_world_size() +def _get_torch_distributed_rank() -> int: + if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_rank"): + return 0 + return torch.distributed.get_rank() + + def is_local_dist_rank_0(): return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0 @@ -3430,9 +3436,9 @@ def save_pretrained( os.makedirs(save_directory, exist_ok=True) save_on_this_rank = is_main_process - if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is not None: + if _is_torch_distributed_initialized() and getattr(self, "device_mesh", None) is not None: # DTensor gather materializes a full plain state dict only on global rank 0. - save_on_this_rank = save_on_this_rank and torch.distributed.get_rank() == 0 + save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0 save_directory_path = os.fspath(save_directory) if push_to_hub and save_on_this_rank: @@ -3509,7 +3515,7 @@ def save_pretrained( model_to_save.config.distributed_config = distributed_config if distributed_checkpoint: - if torch.distributed.is_initialized() and getattr(self, "device_mesh", None) is None: + if _is_torch_distributed_initialized() and getattr(self, "device_mesh", None) is None: raise ValueError( "save_pretrained(distributed_checkpoint=True) requires the model to have been " "initialized with a distributed_config (device_mesh is None)." From c3d75af342c543c1d9d5295ef5f789de164ebd9f Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 16 Jun 2026 04:13:24 +0000 Subject: [PATCH 7/9] fix --- src/transformers/integrations/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index e864a9936141..1c4a619b2ff3 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,6 @@ "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], - "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "gemma_quant": [ "QuantizedEmbedding", "QuantizedLinear", @@ -210,7 +209,6 @@ from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear - from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .gemma_quant import ( QuantizedEmbedding, QuantizedLinear, From a5401678879ec03a3e4c4d3057c53666b7406f6a Mon Sep 17 00:00:00 2001 From: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Date: Tue, 16 Jun 2026 13:30:54 +0900 Subject: [PATCH 8/9] simplify fsdp (#46679) * simplify fsdp * add verify_fsdp_plan * linting * remove tie_weights() before distribute_model. Now use config.tie_word_embeddings for plan expectations, not pre-shard aliasing. --- .gitignore | 2 +- src/transformers/distributed/__init__.py | 4 +- .../distributed/configuration_utils.py | 9 +- src/transformers/distributed/fsdp.py | 377 ++++++------------ src/transformers/distributed/utils.py | 2 +- src/transformers/modeling_utils.py | 34 +- .../models/qwen3/configuration_qwen3.py | 4 +- .../configuration_qwen3_omni_moe.py | 4 +- tests/test_distributed_config.py | 12 +- tests/test_fsdp_mixin.py | 126 ++---- tests/utils/test_modeling_utils.py | 3 +- 11 files changed, 196 insertions(+), 381 deletions(-) diff --git a/.gitignore b/.gitignore index 903efc854eef..de7ec2b0ac2b 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,7 @@ __pycache__/ # C extensions *.so - +checkpoints/ # tests and logs tests/fixtures/cached_*_text.txt logs/ diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index 09f81832c8e0..b109d902d74d 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -19,7 +19,7 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], - "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"], "sharding_utils": [], "tensor_parallel": [ "ALL_PARALLEL_STYLES", @@ -34,7 +34,7 @@ from .configuration_utils import ( DistributedConfig, ) - from .fsdp import is_fsdp_enabled, is_fsdp_managed_module + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan from .tensor_parallel import ( ALL_PARALLEL_STYLES, apply_tensor_parallel, diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 15c6a9e367f2..12bcd3511f61 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -35,8 +35,10 @@ class DistributedConfig: Tensor parallel sharding plan. Leave as `None` to use the model's default plan. fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. - fsdp_plan (`dict`, *optional*): - FSDP wrapping plan. Leave as `None` to wrap each transformer layer + root. + fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): + Whether to enable CPU offloading for FSDP2. + fsdp_mixed_precision (`bool`, *optional*, defaults to `False`): + Whether to enable mixed precision for FSDP2. """ tp_size: int | None = None @@ -44,7 +46,8 @@ class DistributedConfig: enable_sequence_parallel: bool = False enable_expert_parallel: bool = False fsdp_size: int | None = None - fsdp_plan: dict | None = None + fsdp_cpu_offload: bool = False + fsdp_mixed_precision: bool = False def __post_init__(self): if self.tp_size is None and self.fsdp_size is None: diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 64395ff2f531..4483a3b5d744 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -15,7 +15,7 @@ import inspect import os -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any from ..utils import is_torch_available, is_torch_greater_or_equal, logging, strtobool from ..utils.quantization_config import QuantizationMethod @@ -25,12 +25,14 @@ if TYPE_CHECKING: import torch.nn as nn + from .configuration_utils import DistributedConfig + if is_torch_available(): import torch if is_torch_available() and is_torch_greater_or_equal("2.6"): from torch.distributed._composable.fsdp import fully_shard - from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, OffloadPolicy + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy logger = logging.get_logger(__name__) @@ -66,185 +68,117 @@ def is_fsdp_managed_module(module: nn.Module) -> bool: return isinstance(module, FullyShardedDataParallel) -def _get_policy_kwargs(fsdp_plan: dict[str, Any]) -> dict[str, Any]: - """Parse `cpu_offload` / `mixed_precision` flags from the user fsdp_plan into fully_shard kwargs.""" - policy_kwargs = {} - if fsdp_plan.get("cpu_offload"): - policy_kwargs["offload_policy"] = CPUOffloadPolicy() - if fsdp_plan.get("mixed_precision"): - policy_kwargs["mp_policy"] = MixedPrecisionPolicy( +def _get_fsdp_policy_kwargs(distributed_config: DistributedConfig | None) -> dict[str, Any]: + """Build ``fully_shard`` policy kwargs from ``DistributedConfig`` runtime flags.""" + if distributed_config is None: + return {} + + fsdp_policy_kwargs = {} + if distributed_config.fsdp_cpu_offload: + fsdp_policy_kwargs["offload_policy"] = CPUOffloadPolicy() + if distributed_config.fsdp_mixed_precision: + fsdp_policy_kwargs["mp_policy"] = MixedPrecisionPolicy( param_dtype=torch.bfloat16, reduce_dtype=torch.float32, output_dtype=None, ) - return policy_kwargs + return fsdp_policy_kwargs + + +def is_norm_and_head_pair(modules: list[tuple[str, Any]]) -> bool: + """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" + if len(modules) != 2: + return False + module_names = [name for name, _ in modules] + has_final_norm = any(name == "norm" or name.endswith(".norm") for name in module_names) + has_output_head = any( + name in {"lm_head", "embed_tokens"} or name.endswith((".lm_head", ".embed_tokens")) for name in module_names + ) + return has_final_norm and has_output_head -def _parse_manual_plan_entry( - entry: list[str], -) -> tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: +def _resolve_tied_embed_lm_head_plan( + fsdp_plan: dict[str, str], + *, + tie_word_embeddings: bool, +) -> dict[str, str]: """ - Returns: - tuple[bool, MixedPrecisionPolicy | None, CPUOffloadPolicy | None]: - - bool: whether to reshard after forward - - MixedPrecisionPolicy | None: mixed precision policy - - CPUOffloadPolicy | None: cpu offload policy + Rewrite the plan so tied embed/lm_head weights are wrapped once. + Example: + {"model.embed_tokens": "free_full_weight", + "model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "lm_head": "keep_full_weight"} + -> + {"model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "model.embed_tokens": "keep_full_weight"} """ + if not tie_word_embeddings: + return fsdp_plan - if not isinstance(entry, list): - raise ValueError( - f"Manual fsdp_plan values must be a list of strings combining strategy/policies, got {type(entry)}" - ) - items = entry - - strategy: Literal["free_full_weight", "keep_full_weight"] | None = None - offload_policy: CPUOffloadPolicy | None = None - mp_policy: MixedPrecisionPolicy | None = None - - for item in items: - if not isinstance(item, str): - raise ValueError( - f"fsdp_plan option must be a string, got {type(item)}. " - "Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." - ) - token = item.lower() - if token in {"free_full_weight", "keep_full_weight"}: - strategy = token - elif token == "cpu_offload": - offload_policy = CPUOffloadPolicy() - elif token == "mixed_precision": - # TODO(3outeille): add support for different dtypes - mp_policy = MixedPrecisionPolicy( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - output_dtype=torch.bfloat16, - ) - else: - raise ValueError( - "Unknown fsdp_plan option " - f"{item!r}. Supported: 'free_full_weight', 'keep_full_weight', 'cpu_offload', 'mixed_precision'." - ) - - if strategy is None: - strategy = "free_full_weight" - - return strategy != "keep_full_weight", mp_policy, offload_policy - - -def _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): - if pattern in name_to_module: - target = name_to_module[pattern] - if isinstance(target, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential)): - # (ModuleList, ModuleDict, Sequential) don't have a forward() that gets called - - # the model loops over their children directly. So when a pattern matches a - # container, we shard each child instead. - for child_name, child in target.named_children(): - yield f"{pattern}.{child_name}", child - else: - yield pattern, target - return + embed_module_key = next( + (key for key in fsdp_plan if key == "embed_tokens" or key.endswith(".embed_tokens")), + None, + ) + if embed_module_key is None: + return fsdp_plan - # Prefix match: "model.layers" matches "model.layers.0", etc. - for name, module in model.named_modules(): - if name in already_sharded_names or isinstance( - module, (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential) - ): - continue - if name != pattern and not name.startswith(pattern + "."): + adapted_plan = {} + for key, sharding_strategy in fsdp_plan.items(): + if key == embed_module_key: continue - if any( - name.startswith(already_sharded_names_name + ".") for already_sharded_names_name in already_sharded_names - ): - continue - yield name, module - + if key == "lm_head" and sharding_strategy == "keep_full_weight": + adapted_plan[embed_module_key] = sharding_strategy + else: + adapted_plan[key] = sharding_strategy + return adapted_plan -def _get_manual_plan_modules(fsdp_plan: dict[str, Any]) -> dict[str, list[str]]: - modules = fsdp_plan.get("modules") - if not isinstance(modules, dict): - raise ValueError("Manual fsdp_plan must define a 'modules' dict.") - return modules +def expand_fsdp_plan(model, fsdp_plan: dict[str, str]) -> list[tuple[str, nn.Module, str]]: + """Expand plan keys into ``(module_name, module, sharding_strategy)`` shard targets.""" + module_lookup = dict(model.named_modules()) + shard_targets: list[tuple[str, nn.Module, str]] = [] -def is_tail_pair(entries) -> bool: - """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" - if len(entries) != 2: - return False - names = [name for name, _ in entries] - has_norm = any(n == "norm" or n.endswith(".norm") for n in names) - has_head = any(n in {"lm_head", "embed_tokens"} or n.endswith((".lm_head", ".embed_tokens")) for n in names) - return has_norm and has_head - + for plan_key, sharding_strategy in fsdp_plan.items(): + if plan_key in module_lookup: + shard_targets.append((plan_key, module_lookup[plan_key], sharding_strategy)) + continue -def tied_source_path(model) -> str | None: - """Return the dotted path of the input embedding module (the tied source).""" - input_embed = getattr(model, "get_input_embeddings", lambda: None)() - if input_embed is None: - return None - for name, module in model.named_modules(): - if module is input_embed: - return name - return None + for module_name, module in module_lookup.items(): + if replace_layer_number_by_wildcard(module_name) == plan_key: + shard_targets.append((module_name, module, sharding_strategy)) + return shard_targets -def _resolve_plan_key(name_to_module: dict, key: str): - """Resolve a plan key into the matching (name, module) pairs. - Supports exact module names and tp_plan-style wildcards (via - ``replace_layer_number_by_wildcard``). +def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) -> None: """ - if key in name_to_module: - return [(key, name_to_module[key])] - return [(name, mod) for name, mod in name_to_module.items() if replace_layer_number_by_wildcard(name) == key] - + Verify the FSDP plan of the model, log a warning if plan keys were not applied or strategies are invalid. + """ + if not fsdp_plan: + return -def _iter_plan_targets(model, plan, is_weights_tied: bool, tied_source: str | None): - """Yield ``(name, module, strategy)`` for every module the plan applies to. + name_lookup = dict.fromkeys(module_names) + unused_rules: dict[str, str] = {} + invalid_strategies: dict[str, str] = {} - Expands wildcards via ``_resolve_plan_key`` and pre-applies tying rules: - skips the standalone tied-source entry, and rewrites a keep ``"lm_head"`` - entry to the tied source so the shared parameter is wrapped once. - """ - name_to_module = dict(model.named_modules()) - for key, strategy in plan.items(): - if is_weights_tied and key == tied_source: + for key, strategy in fsdp_plan.items(): + if strategy not in {"free_full_weight", "keep_full_weight"}: + invalid_strategies[key] = strategy continue - if is_weights_tied and key == "lm_head" and strategy == "keep_full_weight" and tied_source is not None: - yield tied_source, name_to_module[tied_source], strategy - continue - for name, module in _resolve_plan_key(name_to_module, key): - yield name, module, strategy + if key not in name_lookup and not any(replace_layer_number_by_wildcard(name) == key for name in name_lookup): + unused_rules[key] = strategy + + if invalid_strategies: + logger.warning(f"The following FSDP entries have unknown strategies: {invalid_strategies}") + if unused_rules: + logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}") -def apply_fully_shard_data_parallel( - model, - fsdp_mesh, - fsdp_plan: dict[str, Any] | None, -): +def apply_fully_shard_data_parallel(model, fsdp_mesh): """ Apply FSDP2 (fully_shard) to a model. - - When ``fsdp_plan`` is ``None`` or doesn't contain a ``"modules"`` key, the - model-declared ``model._fsdp_plan`` drives sharding. Policies (`cpu_offload`, - `mixed_precision`) from ``fsdp_plan`` are applied on top. - - When ``fsdp_plan`` has a ``"modules"`` key, the user fully specifies the - layout (manual mode). - - Examples: - # Plan-driven (uses model._fsdp_plan). - fsdp_plan = None - fsdp_plan = {"cpu_offload": True, "mixed_precision": True} - - # Manual override. - fsdp_plan = { - "modules": { - "model.embed_tokens": ["free_full_weight"], - "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], - "model.norm": ["keep_full_weight"], - "lm_head": ["keep_full_weight"], - }, - } """ if not is_torch_available(): raise ImportError("PyTorch is required for FSDP support") @@ -252,108 +186,53 @@ def apply_fully_shard_data_parallel( if not is_torch_greater_or_equal("2.6"): raise OSError("FSDP2 requires torch>=2.6") - if fsdp_plan is None: - fsdp_plan = {} - - input_embed = getattr(model, "get_input_embeddings", lambda: None)() - output_embed = getattr(model, "get_output_embeddings", lambda: None)() - is_weights_tied = ( - input_embed is not None - and output_embed is not None - and hasattr(input_embed, "weight") - and hasattr(output_embed, "weight") - and input_embed.weight is output_embed.weight - ) - - if not isinstance(fsdp_plan, dict): - raise ValueError(f"fsdp_plan must be a dict, got {type(fsdp_plan)}") - is_manual = "modules" in fsdp_plan - - if not is_manual: - policy_kwargs = _get_policy_kwargs(fsdp_plan) - - plan = getattr(model, "_fsdp_plan", None) or {} - if not plan: - raise ValueError( - f"{type(model).__name__} has no `_fsdp_plan` declared. Either set " - "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class, " - "or pass an explicit `fsdp_plan={'modules': {...}}` manual override." - ) - - tied_source = tied_source_path(model) if is_weights_tied else None - keep_buffer: list[tuple[str, Any]] = [] - - for name, module, strategy in _iter_plan_targets(model, plan, is_weights_tied, tied_source): - if strategy == "keep_full_weight": - keep_buffer.append((name, module)) - continue - fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **policy_kwargs) - logger.debug(f"Applied fully_shard to {name} (reshard=True)") - - # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) - # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. - if is_tail_pair(keep_buffer): - keep_names = [n for n, _ in keep_buffer] - keep_modules = [m for _, m in keep_buffer] - fully_shard(keep_modules, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) - logger.debug(f"Grouped tail {keep_names} (reshard=False)") - else: - for name, module in keep_buffer: - fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **policy_kwargs) - logger.debug(f"Applied fully_shard to {name} (reshard=False)") + fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) + if not fsdp_plan: + raise ValueError( + f"{type(model).__name__} does not have a FSDP2 plan declared. Set " + "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class." + ) - # Shard root model - fully_shard(model, mesh=fsdp_mesh, **policy_kwargs) + distributed_config = getattr(model.config, "distributed_config", None) + fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) + tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) - logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(plan)} entries") + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, tie_word_embeddings=tie_word_embeddings) + shard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) + reshard_targets = [] + no_reshard_targets = [] + for module_name, module, sharding_strategy in shard_targets: + if sharding_strategy == "keep_full_weight": + no_reshard_targets.append((module_name, module)) + else: + reshard_targets.append((module_name, module)) + + for module_name, module in reshard_targets: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) + logger.debug(f"Applied fully_shard to {module_name} (reshard=True)") + + # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) + # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. + if is_norm_and_head_pair(no_reshard_targets): + module_names = [name for name, _ in no_reshard_targets] + modules = [module for _, module in no_reshard_targets] + fully_shard(modules, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) + logger.debug(f"Grouped tail {module_names} (reshard=False)") else: - # fsdp_plan = { - # "modules": { - # "model.layers.0.self_attn": ["free_full_weight"], # reshard_after_forward=True - # "model.norm": ["keep_full_weight"], # reshard_after_forward=False - # "model.layers.0.mlp": ["free_full_weight", "cpu_offload", "mixed_precision"], - # }, - # } - - name_to_module = dict(model.named_modules()) - already_sharded_names: set[str] = set() - root_mp_policy = MixedPrecisionPolicy() - root_offload_policy = OffloadPolicy() - - for pattern, entry in _get_manual_plan_modules(fsdp_plan).items(): - reshard, mp_policy, offload_policy = _parse_manual_plan_entry(entry) - if mp_policy is not None: - root_mp_policy = mp_policy - if offload_policy is not None: - root_offload_policy = offload_policy - - for name, module in _iter_manual_plan_targets(model, pattern, name_to_module, already_sharded_names): - if name in already_sharded_names: - continue - shard_kwargs = {"mesh": fsdp_mesh, "reshard_after_forward": reshard} - if mp_policy is not None: - shard_kwargs["mp_policy"] = mp_policy - if offload_policy is not None: - shard_kwargs["offload_policy"] = offload_policy - fully_shard(module, **shard_kwargs) - already_sharded_names.add(name) - logger.debug(f"Applied fully_shard to {name}") - - # Shard root model with the same policies as sub-modules. - # MixedPrecisionPolicy.output_dtype casting happens in post_forward - # for every fully_shard-wrapped module, even with no direct parameters. - fully_shard(model, mesh=fsdp_mesh, mp_policy=root_mp_policy, offload_policy=root_offload_policy) + for module_name, module in no_reshard_targets: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) + logger.debug(f"Applied fully_shard to {module_name} (reshard=False)") + + # Apply FSDP2 to the root module + fully_shard(model, mesh=fsdp_mesh, **fsdp_policy_kwargs) + + logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(fsdp_plan)} entries") # Used by generation code to detect FSDP and enable synced_gpus. model._is_fsdp_managed_module = True - if is_weights_tied and hasattr(model, "tie_weights"): - # Re-tie weights. - # fully_shard replaces nn.Parameter objects (swapping data for DTensor shards), - # which breaks weight tying (e.g. lm_head.weight is no longer embed_tokens.weight). - # Re-tying makes lm_head._parameters["weight"] point to the new DTensor parameter - # so gradients accumulate correctly into a single buffer. + if tie_word_embeddings and hasattr(model, "tie_weights"): model.tie_weights() return model diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index ae159e510a7d..855d5f8175dc 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -149,7 +149,7 @@ def distribute_model(model, distributed_config: DistributedConfig, device_mesh) model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) if "fsdp" in mesh_dim_names: fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh - model = apply_fully_shard_data_parallel(model, fsdp_mesh, distributed_config.fsdp_plan) + model = apply_fully_shard_data_parallel(model, fsdp_mesh) return model diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 2a2bd200e12c..bd8869e9f3f9 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -53,7 +53,7 @@ revert_weight_conversion, ) from .distributed import DistributedConfig -from .distributed.fsdp import is_fsdp_enabled +from .distributed.fsdp import is_fsdp_enabled, verify_fsdp_plan from .distributed.sharding_utils import _dtensor_from_local_like from .distributed.tensor_parallel import ( _get_parameter_tp_plan, @@ -186,6 +186,7 @@ class LoadStateDictConfig: hf_quantizer: HfQuantizer | None = None device_mesh: "DeviceMeshLike | None" = None tp_plan: dict[str, str] | None = None + fsdp_plan: dict[str, str] | None = None weights_only: bool = True weight_mapping: list[WeightConverter | WeightRenaming] | None = None disable_mmap: bool | None = None @@ -1488,6 +1489,10 @@ def tp_plan(self) -> dict[str, str]: return self._ep_plan return self._tp_plan + @property + def fsdp_plan(self) -> dict[str, str]: + return self._fsdp_plan + @property def pp_plan(self) -> dict[str, tuple[str, str]]: return self._pp_plan @@ -4064,16 +4069,17 @@ def from_pretrained( exclusive with `quantization_config` (for now) and `device_map`. When set, accelerate is not used for device placement or dispatch. Launch with `torchrun --nproc_per_node=N script.py`. - Accepts `tp_size`, `tp_plan`, `fsdp_size`, `fsdp_plan`. When a size is specified without a - plan, the plan defaults to the model's predefined behavior: TP uses the model's predefined - tensor parallel sharding plan, FSDP wraps each transformer layer individually with FSDP2 - (`fully_shard`). Both plans also accept a `dict` for manual control: `tp_plan` maps - parameter names to parallel styles (e.g. `{"model.layers.*.self_attn.q_proj": "colwise"}`), - `fsdp_plan` maps module names to wrap (e.g. `{"model.layers.0": {}, "model.layers.1": {}}`). + Accepts `tp_size`, `tp_plan`, `fsdp_size`, `fsdp_cpu_offload`, and + `fsdp_mixed_precision`. When a size is specified without a plan, the plan defaults to + the model's predefined behavior: TP uses the model's predefined tensor parallel + sharding plan, FSDP uses the model's ``base_model_fsdp_plan`` with FSDP2 + (`fully_shard`). `tp_plan` can override the tensor parallel layout (e.g. + `{"model.layers.*.self_attn.q_proj": "colwise"}`). Examples: - TP-only: `DistributedConfig(tp_size=4)` - FSDP-only: `DistributedConfig(fsdp_size=4)` + - FSDP with mixed precision: `DistributedConfig(fsdp_size=4, fsdp_mixed_precision=True)` - 2D parallel: `DistributedConfig(tp_size=2, fsdp_size=2)` on 4 GPUs device_mesh (`torch.distributed.DeviceMesh`, *optional*): A torch device mesh. If not provided would default to world size. Used only for tensor parallel for now. @@ -4384,11 +4390,6 @@ def from_pretrained( weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) if distributed_config is not None: - # Tie weights before sharding so `apply_fully_shard_data_parallel` / - # `apply_tensor_parallel` see the shared-parameter graph and can route tied - # entries (e.g. `lm_head` -> `embed_tokens`) correctly. `_finalize_model_loading` - # re-runs `tie_weights` after the checkpoint is loaded to handle missing-key edge cases. - model.tie_weights() model = distribute_model(model, distributed_config, device_mesh) elif device_map is not None: # Expand device_map if it was passed as a `str`, i.e. `device_map="auto"` @@ -4410,6 +4411,7 @@ def from_pretrained( hf_quantizer=hf_quantizer, device_mesh=device_mesh, tp_plan=active_tp_plan, + fsdp_plan=model.fsdp_plan, weights_only=weights_only, weight_mapping=weight_conversions, use_safetensors=use_safetensors, @@ -4475,10 +4477,12 @@ def _load_pretrained_model( } # Model's definition arriving here is final (TP hooks added, quantized layers replaces) - expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys + expected_tp_plan_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys + expected_fsdp_plan_keys = [name for name, _ in model.named_modules()] if logger.level >= logging.WARNING: - verify_tp_plan(expected_keys, load_config.tp_plan) + verify_tp_plan(expected_tp_plan_keys, load_config.tp_plan) + verify_fsdp_plan(expected_fsdp_plan_keys, load_config.fsdp_plan) # This offload index if for params explicitly on the "disk" in the device_map disk_offload_index = None @@ -4496,7 +4500,7 @@ def _load_pretrained_model( # Warmup cuda to load the weights much faster on devices if load_config.device_map is not None and not is_hqq_or_quark: - expanded_device_map = expand_device_map(load_config.device_map, expected_keys) + expanded_device_map = expand_device_map(load_config.device_map, expected_tp_plan_keys) caching_allocator_warmup(model, expanded_device_map, load_config.hf_quantizer) error_msgs = [] diff --git a/src/transformers/models/qwen3/configuration_qwen3.py b/src/transformers/models/qwen3/configuration_qwen3.py index ff0cd8f63a23..9e573bccbd7f 100644 --- a/src/transformers/models/qwen3/configuration_qwen3.py +++ b/src/transformers/models/qwen3/configuration_qwen3.py @@ -83,8 +83,8 @@ class Qwen3Config(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) - # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # FSDP2 plan. Values are sharding strategies; runtime flags (cpu_offload, mixed_precision) + # live on DistributedConfig. All entries marked `keep_full_weight` are bundled # into a single fully_shard([...]) call at apply time so they share one all-gather. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py index 3c65badeafb4..b18e65f61b74 100644 --- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py @@ -310,8 +310,8 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PreTrainedConfig): "norm": (["hidden_states"], ["hidden_states"]), } - # FSDP2 plan. Values are sharding strategies; policies (cpu_offload, mixed_precision) - # live on DistributedConfig.fsdp_plan. All entries marked `keep_full_weight` are bundled + # FSDP2 plan. Values are sharding strategies; runtime flags (cpu_offload, mixed_precision) + # live on DistributedConfig. All entries marked `keep_full_weight` are bundled # into a single fully_shard([...]) call at apply time so they share one all-gather. base_model_fsdp_plan = { "embed_tokens": "free_full_weight", diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py index 3cf5b4a914db..d6b2b0ab2bd5 100644 --- a/tests/test_distributed_config.py +++ b/tests/test_distributed_config.py @@ -10,7 +10,8 @@ def test_2d_parallelism(self): assert dc.tp_size == 2 assert dc.fsdp_size == 2 assert dc.tp_plan is None - assert dc.fsdp_plan is None + assert dc.fsdp_cpu_offload is False + assert dc.fsdp_mixed_precision is False def test_tp_only_defaults_fsdp_to_1(self): dc = DistributedConfig(tp_size=4) @@ -22,7 +23,8 @@ def test_fsdp_only_defaults_tp_to_1(self): dc = DistributedConfig(fsdp_size=4) assert dc.tp_size == 1 assert dc.fsdp_size == 4 - assert dc.fsdp_plan is None + assert dc.fsdp_cpu_offload is False + assert dc.fsdp_mixed_precision is False assert dc.tp_plan is None def test_empty_config(self): @@ -30,7 +32,8 @@ def test_empty_config(self): assert dc.tp_size is None assert dc.fsdp_size is None assert dc.tp_plan is None - assert dc.fsdp_plan is None + assert dc.fsdp_cpu_offload is False + assert dc.fsdp_mixed_precision is False def test_from_dict(self): dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) @@ -56,7 +59,8 @@ def test_to_dict(self): "enable_sequence_parallel": False, "enable_expert_parallel": False, "fsdp_size": 4, - "fsdp_plan": None, + "fsdp_cpu_offload": False, + "fsdp_mixed_precision": False, } def test_to_dict_is_a_copy(self): diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 88734b501dfe..280ea724f63a 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -46,7 +46,7 @@ from torch.nn.parallel import DistributedDataParallel as DDP from transformers.distributed import DistributedConfig - from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard + from transformers.distributed.fsdp import _resolve_tied_embed_lm_head_plan, expand_fsdp_plan from transformers.distributed.utils import ( gather_full_state_dict, load_optimizer_distributed, @@ -217,63 +217,11 @@ def _gather_ddp_state_dict(model): return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()} -def _resolve_fsdp_plan_paths(model): - """Expand model._fsdp_plan into (paths, strategy) entries. - - Wildcard keys are expanded via ``replace_layer_number_by_wildcard``. When - weights are tied, the standalone embed_tokens entry is skipped and any - ``"lm_head"`` keep entry is rewritten to the tied source path (the keep - group will wrap the shared parameter once). - """ - plan = model._fsdp_plan - - input_embed = model.get_input_embeddings() - output_embed = model.get_output_embeddings() - weights_tied = ( - input_embed is not None - and output_embed is not None - and hasattr(input_embed, "weight") - and hasattr(output_embed, "weight") - and input_embed.weight is output_embed.weight - ) - tied_source = None - if weights_tied: - for name, mod in model.named_modules(): - if mod is input_embed: - tied_source = name - break - - name_to_module = dict(model.named_modules()) - entries: list[tuple[list[str], str]] = [] - for key, strategy in plan.items(): - if weights_tied and key == tied_source: - continue - if weights_tied and key == "lm_head" and strategy == "keep_full_weight": - entries.append(([tied_source], strategy)) - continue - - if key in name_to_module: - entries.append(([key], strategy)) - continue - matched = [name for name in name_to_module if replace_layer_number_by_wildcard(name) == key] - if matched: - entries.append((matched, strategy)) - return entries - - -def _build_manual_fsdp_plan(config, device, policy_options=None): - """Build a manual FSDP2 plan by expanding model._fsdp_plan.""" - policy_options = policy_options or [] - set_seed(SEED) - model = AutoModelForCausalLM.from_config(config).to(device) - - module_plan: dict[str, list[str]] = {} - for paths, strategy in _resolve_fsdp_plan_paths(model): - for path in paths: - module_plan[path] = [strategy, *policy_options] - - del model - return {"modules": module_plan} +def _fsdp_target_module_names(model) -> set[str]: + """Return module names that should receive a ``fully_shard`` call (excluding root).""" + tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, tie_word_embeddings=tie_word_embeddings) + return {name for name, _, _ in expand_fsdp_plan(model, adapted_fsdp_plan)} def _save_init_pretrained(rank, config, dtype): @@ -347,11 +295,16 @@ def train_fsdp2( dtype, init_model_dir, checkpoint_step, - fsdp_plan, + fsdp_cpu_offload=False, + fsdp_mixed_precision=False, ): # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save _set_determinism(SEED) - distributed_config = DistributedConfig(fsdp_size=dist.get_world_size(), fsdp_plan=fsdp_plan) + distributed_config = DistributedConfig( + fsdp_size=dist.get_world_size(), + fsdp_cpu_offload=fsdp_cpu_offload, + fsdp_mixed_precision=fsdp_mixed_precision, + ) pre_ckpt_model = AutoModelForCausalLM.from_pretrained( init_model_dir, torch_dtype=dtype, distributed_config=distributed_config ) @@ -481,7 +434,7 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict): def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings): """ - Verify that apply_fully_shard_data_parallel(fsdp_plan=None) wraps exactly the right modules. + Verify that apply_fully_shard_data_parallel wraps exactly the right modules. Expected FSDP targets: UNTIED TIED @@ -510,9 +463,7 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor # Expected FSDP targets come from model._fsdp_plan: every resolved path gets a # fully_shard call (keep_full_weight entries are bundled into one group, but each # member still appears as an FSDP-wrapped module in named_modules), plus the root. - expected_targets = {""} - for paths, _strategy in _resolve_fsdp_plan_paths(model): - expected_targets.update(paths) + expected_targets = {""} | _fsdp_target_module_names(model) actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")} @@ -534,9 +485,9 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor def _test_fsdp2_plan_vs_ddp_impl( - rank, config_class, config_dict, tie_word_embeddings, plan_mode, policy_options=None, dtype=None + rank, config_class, config_dict, tie_word_embeddings, policy_options=None, dtype=None ): - """Validate DDP-vs-FSDP2 trace matching for either auto or manual plan mode.""" + """Validate DDP-vs-FSDP2 trace matching using the model's declared FSDP plan.""" init_test_logger() if dtype is None: @@ -551,17 +502,9 @@ def _test_fsdp2_plan_vs_ddp_impl( config = config_class.from_dict(config_dict) config.tie_word_embeddings = tie_word_embeddings - if plan_mode == "auto": - fsdp_plan = { - "cpu_offload": "cpu_offload" in policy_options, - "mixed_precision": "mixed_precision" in policy_options, - } - test_label = f"FSDP2(auto{'+policies' if policy_options else ''})" - elif plan_mode == "manual": - fsdp_plan = _build_manual_fsdp_plan(config, device, policy_options=policy_options) - test_label = f"FSDP2(manual{'+policies' if policy_options else ''})" - else: - raise ValueError(f"Unsupported plan_mode '{plan_mode}'. Expected 'auto' or 'manual'.") + fsdp_cpu_offload = "cpu_offload" in policy_options + fsdp_mixed_precision = "mixed_precision" in policy_options + test_label = f"FSDP2{'+policies' if policy_options else ''}" checkpoint_step = NUM_STEPS // 2 init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype) @@ -576,7 +519,8 @@ def _test_fsdp2_plan_vs_ddp_impl( dtype, init_model_dir=init_model_dir, checkpoint_step=checkpoint_step, - fsdp_plan=fsdp_plan, + fsdp_cpu_offload=fsdp_cpu_offload, + fsdp_mixed_precision=fsdp_mixed_precision, ) finally: if rank == 0 and init_tmpdir_obj is not None: @@ -757,28 +701,10 @@ def test_fsdp2_save_load(self): @is_fsdp_test @require_fsdp - def test_fsdp2_auto_plan_vs_ddp_untied(self): - self._run_fsdp2_distributed_test( - "test_fsdp2_auto_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "auto" - ) - - @is_fsdp_test - @require_fsdp - def test_fsdp2_auto_plan_vs_ddp_tied(self): - self._run_fsdp2_distributed_test( - "test_fsdp2_auto_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "auto" - ) + def test_fsdp2_plan_vs_ddp_untied(self): + self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False) @is_fsdp_test @require_fsdp - def test_fsdp2_manual_plan_vs_ddp_untied(self): - self._run_fsdp2_distributed_test( - "test_fsdp2_manual_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False, "manual" - ) - - @is_fsdp_test - @require_fsdp - def test_fsdp2_manual_plan_vs_ddp_tied(self): - self._run_fsdp2_distributed_test( - "test_fsdp2_manual_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True, "manual" - ) + def test_fsdp2_plan_vs_ddp_tied(self): + self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True) diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index e620edb8060d..a83bfdf20a16 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -449,10 +449,9 @@ def test_model_from_pretrained_fsdp_distributes_before_loading(self): fake_mesh.ndim = 1 fake_mesh.device_type = "cpu" - def fake_apply_fsdp(model, fsdp_mesh, fsdp_plan): + def fake_apply_fsdp(model, fsdp_mesh): call_order.append("distribute") self.assertIs(fsdp_mesh, fake_mesh) - self.assertIsNone(fsdp_plan) return model def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None): From 30972993f8a72aee87340f4f21b04c2c2cc72bae Mon Sep 17 00:00:00 2001 From: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Date: Wed, 17 Jun 2026 11:43:05 +0900 Subject: [PATCH 9/9] Tp param level (#46290) * [distributed] Add param-level MoE TP/EP styles and ep_router Decompose MoE tensor/expert parallelism per review feedback: weight sharding is declared per-parameter, while the experts module entry stays forward-comm only. - MoEParamShard: parameter-only style wrapping named expert weights as DTensor placeholders (no forward hook). grouped_gemm shards the expert dim and updates module.num_experts to the per-rank local count. - Register grouped_gemm (Shard(0)), moe_gate_up_colwise (_StridedShard(-2)), moe_gate_up_colwise_alt (_StridedShard(-1)), moe_down_rowwise (Shard(-1)). - EpRouterParallel (ep_router): forward-only slicing of router outputs to local experts, ported from the original RouterParallel (#39501). - moe_experts_allreduce is now forward-comm only: strip the baked shard_plan and drop the now-dead shard_plan ctor arg / _moe_shard_plan / shard_parameters override from MoEExpertsParallel; skip _AllReduceBackward on routing weights under EP. - verify_tp_plan: treat moe_experts_allreduce / ep_router as forward-only. * [distributed] Add param-level apply pass to apply_tensor_parallel Run tensor parallelism in two passes: - Pass 1 (param-level): walk named_parameters() and, for styles in PARAM_ONLY_STYLES (grouped_gemm, moe_gate_up_colwise[_alt], moe_down_rowwise), shard the parameter directly via shard_parameters(). No forward hook. - Pass 2 (module-level): the existing named_modules() loop for forward hooks, now skipping PARAM_ONLY_STYLES. Param sharding runs first so module forward hooks (moe_experts_allreduce) see the already-sharded DTensor params. Also wire the EP-plan fallback so enable_expert_parallel uses model._ep_plan when no explicit plan is passed. * [distributed] Add MoE TP/EP plan tests and a two-sided sharding assertion - New tests/distributed/test_moe_tensor_parallel_plan.py: plan resolution, placement expectations for grouped_gemm / moe_gate_up_colwise[_alt] / moe_down_rowwise, gloo distributed integration (EP Shard(0), TP _StridedShard(-2)+Shard(-1), ep_router slicing), and a registry guard that moe_experts_allreduce carries no baked shard plan. - _verify_tp_sharding: add a two-sided check asserting that every parameter whose plan entry is a weight-sharding style actually comes back as a non-replicate DTensor. The prior check only validated params that happened to be sharded, so a style that gracefully degrades to replicated when unsharded (e.g. MoEExpertsParallel) could pass output-equality while silently running unparallelized. * [distributed] Migrate MoE configs to decomposed TP/SP expert plans For every TP/SP plan that sharded experts, declare per-parameter entries: "layers.*.mlp.experts.gate_up_proj": "moe_gate_up_colwise" "layers.*.mlp.experts.down_proj": "moe_down_rowwise" while keeping the forward-only "layers.*.mlp.experts": "moe_experts_allreduce". This matches the now-empty moe_experts_allreduce shard_plan; sharding is declared in config at parameter granularity. EP plans already used "grouped_gemm" and are unchanged. hy_v3 and laguna previously used "packed_colwise" / "rowwise_allreduce" on the 3D expert *parameters*; those styles are module-level and were silently no-ops on params (the bundled shard_plan did the work). They now use the param-level moe_gate_up_colwise / moe_down_rowwise like every other MoE model. Edited modular files where they own the plan literal; generated configs and inherited plans (e.g. from qwen3_moe) propagated via modular conversion. * [distributed] Document decomposed MoE TP/EP plans - expert_parallelism.md: describe the param-level decomposition (grouped_gemm, ep_router, moe_experts_allreduce) instead of the removed GroupedGemmParallel class, and note the TP equivalents (moe_gate_up_colwise / moe_down_rowwise). - weightconverter.md: note that fused expert weights are sharded at parameter granularity by the parallel plan. * [distributed] Rename MoE intra-expert TP styles to moe_tp_* Rename registry and plan entries so TP-on-expert sharding is distinct from EP (grouped_gemm) and dense packed_colwise: moe_gate_up_colwise -> moe_tp_gate_up_colwise, moe_down_rowwise -> moe_tp_down_rowwise. Drop unused moe_tp_gate_up_colwise_alt (GPT-OSS-style layouts stay EP-only). * handle sparse and dense sp plan for qwen3_moe * better tests coverage for sp & ep * linting * uniformize TP Api to avoid confusion with torch native ops * inline tp * rename * cleaning * inline * cleaning * cleaning * linting * fix ci ep_backward * linting * remove flag expert parallel * fix * add tp plan + ep_plan * revert doc * fix install_forward * linting * add moe identity back * no need aymore * update tp_plan for ernie4_5_vl_moe * sp + ep training / tp + ep inference (#46292) * [distributed] Add resolve_parallel_plan merge helper Compose SP/TP dense recipes with an optional EP overlay and strip intra-expert moe_tp_* when expert parallelism is enabled. Add unit tests for training (SP+EP), inference (TP+EP), and TP-only paths. * [distributed] Wire resolve_parallel_plan into apply_tensor_parallel Replace exclusive SP|EP|TP plan selection with merged plans when tp_plan is unset. Add distributed test for TP+EP merged expert sharding. * [distributed] Use merged plan in tp_plan property and load path Expose resolve_parallel_plan via PreTrainedModel.tp_plan and set active_tp_plan during from_pretrained so checkpoint sharding matches the applied layout. * [distributed] Drop intra-expert moe_tp_* from MoE SP plans Expert weight TP under sequence parallelism comes from the EP overlay (grouped_gemm) when enable_expert_parallel is set; keep moe_tp_* only in base_model_tp_plan for TP-only MoE. * [distributed] Document SP+EP and TP+EP flag combinations Update expert_parallelism guide and DistributedConfig docs for merged plans. Export resolve_parallel_plan and extend resolve-plan tests for trimmed SP sources. * refactor merging plans * add test sp_ep and tp_ep * extend verify_tp_plan to verify_tp_sp_ep_plan * add ep_plan to mixtral and olmoe * cleaning _accumulate_local_param_grad (#46394) * remove _accumulate_local_param_grad * comments * linting * fix * clean _accumulate_local_param_grad * linting * cleaning * cleaning * fix mellun test because of bug in parsing sp_ep plan with regex * aea * Add select_parallel_plan and explicit combo plan config fields Introduce base_model_tp_ep_plan / base_model_sp_ep_plan on PreTrainedConfig, select_parallel_plan() with legacy resolve_parallel_plan fallback, and wire apply_tensor_parallel to use the selector. Model post_init tracks _tp_ep_plan and _sp_ep_plan for composite models. * Add base_model_tp_ep_plan and base_model_sp_ep_plan for Mixtral and Qwen3-MoE Define complete inference TP+EP and training SP+EP plans on the pilot MoE configs. Qwen3-MoE expands per-layer entries in _update_parallel_plans. Add plan_utils and golden tests against legacy resolve_parallel_plan merge. * Add explicit tp_ep / sp_ep plans for remaining MoE models Populate combo plans via init_combo_plans() at config init time for MoE configs that still use split tp/sp/ep recipes. Dynamic configs call it after _update_sp_plan(); modular sources updated for generated configuration files. * Remove resolve_parallel_plan and use explicit combo plan selection Delete runtime plan merging; select_parallel_plan now requires a complete combo dict and raises when missing. apply_tensor_parallel uses DistributedConfig flags directly for SP/EP behavior. Drop model._ep_plan aggregation; load-time verification checks the active plan only. Refresh combo plans after MXFP4 quantizer patches. * Sync modular MoE configs and update expert parallelism docs Propagate init_combo_plans from modular sources to generated configuration files and document select_parallel_plan combo lookup in expert_parallelism.md. * Refactor select_parallel_plan flag lookup for readability Use explicit if/elif branches for the SP/EP flag matrix and derive config_attr from plan_attr instead of parallel lookup dicts. * Write explicit combo parallel plans in MoE configs and remove plan_utils Define base_model_tp_ep_plan and base_model_sp_ep_plan directly in each MoE configuration (or via config-time _update_parallel_plans for dynamic models). Delete plan_utils.py and all init_combo_plans / refresh_combo_plans usage. * Add lm_head entries to _tp_ep_plan and _sp_ep_plan on CausalLM classes Explicit combo plan selection no longer merges _sp_plan with _ep_plan, so head-level lm_head rules must live on _tp_ep_plan/_sp_ep_plan directly. Fixes SP+EP training loss shape mismatch under sequence parallelism. * cleaning * cleaning * cleaning * cleaning * linting * add verify tp and fsdp pla aeaea * revert doc * cleaning * check-repository-consistency * linting --- src/transformers/configuration_utils.py | 11 +- src/transformers/distributed/__init__.py | 2 + .../distributed/configuration_utils.py | 7 +- .../distributed/tensor_parallel.py | 754 ++++++++++++------ src/transformers/distributed/utils.py | 2 +- src/transformers/integrations/moe.py | 6 + src/transformers/modeling_utils.py | 29 +- .../models/afmoe/modeling_afmoe.py | 2 + .../models/afmoe/modular_afmoe.py | 2 + .../models/apertus/modeling_apertus.py | 2 + .../models/arcee/modeling_arcee.py | 2 + src/transformers/models/aria/modeling_aria.py | 2 + .../models/bamba/modeling_bamba.py | 2 + .../models/bitnet/modeling_bitnet.py | 2 + .../models/bitnet/modular_bitnet.py | 2 + .../models/cohere/modeling_cohere.py | 2 + .../models/cohere2/modeling_cohere2.py | 2 + .../cohere2_moe/configuration_cohere2_moe.py | 2 + .../cohere2_moe/modeling_cohere2_moe.py | 2 + src/transformers/models/csm/modeling_csm.py | 2 + src/transformers/models/csm/modular_csm.py | 2 + src/transformers/models/cwm/modeling_cwm.py | 2 + src/transformers/models/dbrx/modeling_dbrx.py | 2 + src/transformers/models/dbrx/modular_dbrx.py | 2 + .../deepseek_v2/configuration_deepseek_v2.py | 4 + .../deepseek_v2/modeling_deepseek_v2.py | 2 + .../models/deepseek_v2/modular_deepseek_v2.py | 4 + .../deepseek_v3/configuration_deepseek_v3.py | 2 + .../deepseek_v3/modeling_deepseek_v3.py | 2 + .../deepseek_v32/modular_deepseek_v32.py | 2 + .../deepseek_v4/modeling_deepseek_v4.py | 2 + .../models/diffllama/modeling_diffllama.py | 2 + .../modular_diffusion_gemma.py | 1 + src/transformers/models/doge/modeling_doge.py | 2 + .../models/dots1/configuration_dots1.py | 2 + .../models/dots1/modeling_dots1.py | 2 + .../models/dots1/modular_dots1.py | 2 + src/transformers/models/emu3/modeling_emu3.py | 2 + .../models/ernie4_5/modeling_ernie4_5.py | 2 + .../configuration_ernie4_5_moe.py | 2 + .../ernie4_5_moe/modeling_ernie4_5_moe.py | 2 + .../configuration_ernie4_5_vl_moe.py | 6 + .../modular_ernie4_5_vl_moe.py | 6 + .../models/eurobert/modeling_eurobert.py | 2 + .../models/eurobert/modular_eurobert.py | 2 + .../models/exaone4/modeling_exaone4.py | 2 + .../exaone_moe/configuration_exaone_moe.py | 8 +- .../models/exaone_moe/modeling_exaone_moe.py | 2 + .../models/exaone_moe/modular_exaone_moe.py | 16 + .../models/falcon_h1/modeling_falcon_h1.py | 2 + .../flex_olmo/configuration_flex_olmo.py | 19 + .../models/flex_olmo/modeling_flex_olmo.py | 2 + .../models/flex_olmo/modular_flex_olmo.py | 19 + .../models/gemma/modeling_gemma.py | 2 + .../models/gemma2/modeling_gemma2.py | 2 + .../models/gemma3/modeling_gemma3.py | 2 + .../models/gemma3n/modeling_gemma3n.py | 2 + .../models/gemma4/configuration_gemma4.py | 12 +- .../models/gemma4/modeling_gemma4.py | 2 + .../modeling_gemma4_assistant.py | 1 + .../gemma4_unified/modular_gemma4_unified.py | 3 + .../modular_gemma4_unified_assistant.py | 1 + src/transformers/models/glm/modeling_glm.py | 2 + src/transformers/models/glm4/modeling_glm4.py | 2 + .../models/glm4_moe/configuration_glm4_moe.py | 2 + .../models/glm4_moe/modeling_glm4_moe.py | 2 + .../models/glm4_moe/modular_glm4_moe.py | 2 + .../configuration_glm4_moe_lite.py | 2 + .../glm4_moe_lite/modeling_glm4_moe_lite.py | 2 + .../glm4_moe_lite/modular_glm4_moe_lite.py | 2 + .../glm4v_moe/configuration_glm4v_moe.py | 6 + .../models/glm4v_moe/modular_glm4v_moe.py | 6 + .../glm_moe_dsa/configuration_glm_moe_dsa.py | 2 + .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 2 + .../models/gpt_neox/modeling_gpt_neox.py | 1 + .../models/gpt_neox/modular_gpt_neox.py | 1 + .../models/gpt_oss/modeling_gpt_oss.py | 2 + .../models/granite/modeling_granite.py | 2 + .../models/granitemoe/modeling_granitemoe.py | 2 + .../modeling_granitemoehybrid.py | 2 + .../modeling_granitemoeshared.py | 2 + .../models/helium/modeling_helium.py | 2 + .../models/hrm_text/modeling_hrm_text.py | 2 + .../modeling_hunyuan_v1_dense.py | 2 + .../hunyuan_v1_moe/modeling_hunyuan_v1_moe.py | 2 + .../models/hy_v3/configuration_hy_v3.py | 4 +- .../models/hy_v3/modeling_hy_v3.py | 2 + .../models/hy_v3/modular_hy_v3.py | 4 +- .../hyperclovax/modeling_hyperclovax.py | 2 + .../models/jais2/modeling_jais2.py | 2 + .../models/jamba/modeling_jamba.py | 2 + .../modeling_kyutai_speech_to_text.py | 2 + .../models/laguna/configuration_laguna.py | 30 +- .../models/laguna/modeling_laguna.py | 2 + .../models/laguna/modular_laguna.py | 23 +- src/transformers/models/lfm2/modeling_lfm2.py | 2 + .../models/lfm2_moe/modeling_lfm2_moe.py | 2 + .../models/llama/modeling_llama.py | 2 + .../models/llama4/configuration_llama4.py | 17 + .../models/llama4/modeling_llama4.py | 2 + .../configuration_longcat_flash.py | 4 + .../longcat_flash/modeling_longcat_flash.py | 2 + .../models/mellum/configuration_mellum.py | 94 ++- .../models/mellum/modeling_mellum.py | 2 + .../models/mellum/modular_mellum.py | 92 ++- .../models/minimax/configuration_minimax.py | 2 + .../models/minimax/modeling_minimax.py | 2 + .../models/minimax/modular_minimax.py | 2 + .../minimax_m2/configuration_minimax_m2.py | 4 + .../models/minimax_m2/modeling_minimax_m2.py | 2 + .../models/minimax_m2/modular_minimax_m2.py | 4 + .../minimax_m3_vl/modular_minimax_m3_vl.py | 2 + .../models/ministral/modeling_ministral.py | 2 + .../models/ministral3/modeling_ministral3.py | 2 + .../models/mistral/modeling_mistral.py | 2 + .../models/mistral4/configuration_mistral4.py | 2 + .../models/mistral4/modeling_mistral4.py | 2 + .../models/mixtral/configuration_mixtral.py | 39 + .../models/mixtral/modeling_mixtral.py | 2 + .../models/nanochat/modeling_nanochat.py | 2 + .../models/nanochat/modular_nanochat.py | 2 + src/transformers/models/olmo/modeling_olmo.py | 2 + .../models/olmo2/modeling_olmo2.py | 2 + .../models/olmo3/modeling_olmo3.py | 2 + .../olmo_hybrid/modeling_olmo_hybrid.py | 2 + .../models/olmoe/configuration_olmoe.py | 37 + .../models/olmoe/modeling_olmoe.py | 2 + src/transformers/models/phi/modeling_phi.py | 2 + src/transformers/models/phi3/modeling_phi3.py | 2 + .../modeling_phi4_multimodal.py | 2 + .../models/phimoe/modeling_phimoe.py | 2 + src/transformers/models/pi0/modeling_pi0.py | 1 + src/transformers/models/pi0/modular_pi0.py | 1 + .../models/qwen2/modeling_qwen2.py | 2 + .../qwen2_moe/configuration_qwen2_moe.py | 23 + .../models/qwen2_moe/modeling_qwen2_moe.py | 2 + .../models/qwen2_moe/modular_qwen2_moe.py | 2 + .../models/qwen3/modeling_qwen3.py | 2 + .../models/qwen3_5/modeling_qwen3_5.py | 2 + .../qwen3_5_moe/configuration_qwen3_5_moe.py | 2 + .../qwen3_5_moe/modeling_qwen3_5_moe.py | 4 + .../models/qwen3_5_moe/modular_qwen3_5_moe.py | 4 + .../qwen3_moe/configuration_qwen3_moe.py | 118 +++ .../models/qwen3_moe/modeling_qwen3_moe.py | 2 + .../qwen3_next/configuration_qwen3_next.py | 2 + .../models/qwen3_next/modeling_qwen3_next.py | 2 + .../configuration_qwen3_omni_moe.py | 119 +++ .../qwen3_omni_moe/modeling_qwen3_omni_moe.py | 4 + .../qwen3_omni_moe/modular_qwen3_omni_moe.py | 42 +- .../configuration_qwen3_vl_moe.py | 117 +++ .../qwen3_vl_moe/modular_qwen3_vl_moe.py | 38 +- .../models/seed_oss/modeling_seed_oss.py | 2 + .../models/smollm3/modeling_smollm3.py | 2 + .../solar_open/configuration_solar_open.py | 4 + .../models/solar_open/modeling_solar_open.py | 2 + .../models/solar_open/modular_solar_open.py | 4 + .../models/starcoder2/modeling_starcoder2.py | 2 + .../models/t5gemma/modeling_t5gemma.py | 1 + .../models/t5gemma/modular_t5gemma.py | 1 + .../models/t5gemma2/modeling_t5gemma2.py | 1 + .../models/t5gemma2/modular_t5gemma2.py | 1 + .../models/vaultgemma/modeling_vaultgemma.py | 2 + .../models/youtu/modeling_youtu.py | 2 + tests/models/gpt_oss/test_modeling_gpt_oss.py | 1 + .../models/hrm_text/test_modeling_hrm_text.py | 8 + .../qwen3_moe/test_modeling_qwen3_moe.py | 9 + tests/test_tensor_parallel_mixin.py | 650 ++++++++------- 167 files changed, 2088 insertions(+), 560 deletions(-) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index c4ff8c46c5b6..4b8ec290c818 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -150,9 +150,12 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): - **base_model_tp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a tensor parallel plan applied to the sub-module when `model.tensor_parallel` is called. - **base_model_sp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a sequence - parallel plan, used in place of `base_model_tp_plan` when `distributed_config.enable_sequence_parallel` is set. - Same key/value shape as the TP plan; values are style names registered in `ALL_PARALLEL_STYLES` - (e.g. `"vocab_reduce_scatter"`, `"rowwise_reduce_scatter"`, `"activation"`, `"module_allgather"`). + parallel plan, used when `distributed_config.enable_sequence_parallel=True` and + `enable_expert_parallel=False`. Same key/value shape as the TP plan. + - **base_model_tp_ep_plan** (`dict[str, Any]`) -- Complete plan for inference TP + expert parallel + (`enable_sequence_parallel=False`, `enable_expert_parallel=True`). + - **base_model_sp_ep_plan** (`dict[str, Any]`) -- Complete plan for training SP + expert parallel + (`enable_sequence_parallel=True`, `enable_expert_parallel=True`). - **base_model_fsdp_plan** (`dict[Any, str]`) -- A dict that maps sub-modules of a base model to an FSDP2 sharding strategy (e.g. `"free_full_weight"` / `"keep_full_weight"`). Keys can be wildcard module paths (e.g. `"layers.*"`) or tuples of paths (grouped into a single `fully_shard` call). @@ -228,6 +231,8 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): attribute_map: ClassVar[dict[str, str]] = {} base_model_tp_plan: ClassVar[dict[str, Any] | None] = None base_model_sp_plan: ClassVar[dict[str, Any] | None] = None + base_model_tp_ep_plan: ClassVar[dict[str, Any] | None] = None + base_model_sp_ep_plan: ClassVar[dict[str, Any] | None] = None base_model_fsdp_plan: ClassVar[dict[Any, str] | None] = None base_model_pp_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None base_model_ep_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index b109d902d74d..8d5570b4abbb 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -24,6 +24,7 @@ "tensor_parallel": [ "ALL_PARALLEL_STYLES", "apply_tensor_parallel", + "select_parallel_plan", "verify_tp_plan", ], "utils": [], @@ -38,6 +39,7 @@ from .tensor_parallel import ( ALL_PARALLEL_STYLES, apply_tensor_parallel, + select_parallel_plan, verify_tp_plan, ) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 12bcd3511f61..05b63e070a6a 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -32,7 +32,12 @@ class DistributedConfig: tp_size (`int`, *optional*): Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1. tp_plan (`dict`, *optional*): - Tensor parallel sharding plan. Leave as `None` to use the model's default plan. + Tensor parallel sharding plan. Leave as `None` to select the model's SP/TP and EP plan + (see ``select_parallel_plan``). Set explicitly to override. + enable_sequence_parallel (`bool`, *optional*, defaults to `False`): + Select ``base_model_sp_plan`` (dense-only) or ``base_model_sp_ep_plan`` (with EP). + enable_expert_parallel (`bool`, *optional*, defaults to `False`): + Select ``base_model_tp_ep_plan`` or ``base_model_sp_ep_plan`` when combined with SP flag. fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 458da9787884..60a8e4a85c70 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -14,6 +14,7 @@ from __future__ import annotations import contextlib +import functools import re from ..utils import logging @@ -27,15 +28,10 @@ if is_torch_available() and is_torch_greater_or_equal("2.5"): import torch.distributed as dist from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor - from torch.distributed.tensor.parallel import ( - ColwiseParallel, - PrepareModuleInput, - RowwiseParallel, - SequenceParallel, - ) - from torch.distributed.tensor.parallel.style import ParallelStyle from torch.distributed.tensor.placement_types import _StridedShard + from .sharding_utils import _find_strided_shard_placement_from_fused_params + # Cache this result has it's a C FFI call which can be pretty time-consuming _torch_distributed_available = torch.distributed.is_available() @@ -63,36 +59,84 @@ def _get_parameter_tp_plan(parameter_name: str, tp_plan: dict[str, str], is_weig The `is_weight` is important because for weights, we want to support `.weights` and `.bias` cases seamlessly! but not parent classes for `post_init` calls """ + # Lookup order (most -> least specific): + # 1. exact param key -> layers.1.mlp.experts.gate_up_proj + if parameter_name in tp_plan: + return tp_plan[parameter_name] generic_param_name = replace_layer_number_by_wildcard(parameter_name) + # 2. wildcard param key -> layers.*.mlp.experts.gate_up_proj if generic_param_name in tp_plan: return tp_plan[generic_param_name] + # 3. parent module (indexed) -> layers.1.mlp.experts + elif is_weight and "." in parameter_name and (module_name := parameter_name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + # 4. parent module (wildcard) -> layers.*.mlp.experts elif is_weight and "." in generic_param_name and (module_name := generic_param_name.rsplit(".", 1)[0]) in tp_plan: return tp_plan[module_name] return None -# ============================================================================= -# High-Level API Functions -# ============================================================================= +class TensorParallelLayer: + def shard_param(self, module, param, mesh): + """Wrap ONE parameter as a DTensor placeholder. Default: no-op.""" + pass + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + return args, kwargs + def context_around_forward(self, module): + return contextlib.nullcontext() + + def transform_output_post_forward(self, module, output, mesh): + return output + + def install_forward(self, module, mesh): + """Install pre / around / post transforms by replacing module.forward.""" + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh) + with self.context_around_forward(module): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh) + + module.forward = tp_forward + return module + + +def _matching_weight_plan_key(generic_key: str, weight_plan: dict[str, str]) -> str | None: + """Return the weight_plan entry that covers a wildcarded parameter name, or None.""" + if generic_key in weight_plan: + return generic_key + # To avoid stripping .gate_up_proj for MoE + if generic_key.endswith(".weight"): + param_name = generic_key.removesuffix(".weight") + elif generic_key.endswith(".bias"): + param_name = generic_key.removesuffix(".bias") + else: + param_name = generic_key + if param_name != generic_key and param_name in weight_plan: + return param_name + if "." in param_name and (parent := param_name.rsplit(".", 1)[0]) in weight_plan: + return parent + return None -def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): - """ - Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules that were not applied. - Only weight-sharding rules (colwise, rowwise, vocab, moe_experts) are checked. - Module/activation entries (e.g. PrepareModuleInput, SequenceParallel) set up - communication hooks on modules, not weight sharding, so they are excluded. +def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None) -> None: """ + Verify the TP plan of the model, log a warning if the layers that were not sharded and the rules + that were not applied. + Only weight-sharding rules (colwise, rowwise, …) that override shard_param are checked. Comm-only entries + (activation, ep_router, moe_experts_allreduce, …) are excluded. + """ if tp_plan is None: return - # Filter out module-level comm hooks — they don't shard weights. - # Plan values are registry names; entries beginning with "activation" or "module" - # configure communication hooks rather than parameter sharding. weight_plan = { - k: v for k, v in tp_plan.items() if not (v == "activation" or v.startswith(("activation_", "module_"))) + k: v + for k, v in tp_plan.items() + if v in ALL_PARALLEL_STYLES and type(ALL_PARALLEL_STYLES[v]).shard_param is not TensorParallelLayer.shard_param } generic_keys = {replace_layer_number_by_wildcard(key) for key in expected_keys} @@ -100,14 +144,8 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): unused_rules = weight_plan.copy() for key in generic_keys: - param_name = key.rsplit(".", 1)[0] if "." in key else key - generic_param_name = re.sub(r"\d+", "*", param_name) - - if generic_param_name in weight_plan: - unused_rules.pop(generic_param_name, None) - unsharded_layers.discard(key) - elif "." in generic_param_name and (parent_param_name := generic_param_name.rsplit(".", 1)[0]) in weight_plan: - unused_rules.pop(parent_param_name, None) + if (matched_key := _matching_weight_plan_key(key, weight_plan)) is not None: + unused_rules.pop(matched_key, None) unsharded_layers.discard(key) if len(unused_rules) > 0: @@ -116,70 +154,132 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None): logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}") -class CustomParallelStyle(ParallelStyle): - """Base class for custom TP styles (MoE, packed linear, etc.). +class ColwiseParallel(TensorParallelLayer): + """Column-wise: weight & bias → Shard(0) (Embedding: Shard(1)); input replicated, output Shard(-1).""" + + def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): + self.input_layouts = input_layouts or Replicate() + self.output_layouts = output_layouts if output_layouts is not None else Shard(-1) + self.use_local_output = use_local_output + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(0) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if x.placements != (Replicate(),): + x = x.redistribute(placements=[Replicate()], async_op=True) + return (x,) + args[1:], kwargs # stay DTensor into F.linear - Apply runs in two phases (subclasses should not override _apply): + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + return output + if output.placements != (self.output_layouts,): + output = output.redistribute(placements=[self.output_layouts], async_op=True) + return output.to_local() if self.use_local_output else output - 1. shard_parameters(module, mesh) — meta-time DTensor placeholders (optional) - 2. _install_forward(module, mesh) — pre / around / post forward transforms + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"output_layouts={self.output_layouts}, use_local_output={self.use_local_output})" + ) - Unlike PyTorch registry entries (ColwiseParallel, SequenceParallel, …) which use - distribute_module or forward hooks via their own _apply, this base wraps - module.forward to support *args, **kwargs and an optional context manager. - Param wrapping runs on meta (the model is on meta when apply_tensor_parallel - is invoked); distribute_tensor on meta builds metadata only — no collective. - Real data flows in later, async, via DtensorShardOperation during load. +class RowwiseParallel(TensorParallelLayer): + """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)). - Override what you need: - - shard_parameters(module, mesh) — wrap params as DTensor placeholders - - transform_inputs_pre_forward(module, args, kwargs, mesh) → (args, kwargs) - - context_around_forward(module) → context manager wrapping the call - - transform_output_post_forward(module, output, mesh) → output + Linear input is sharded on the last dim; Embedding input is replicated. The module + forward produces a Partial output which the boundary redistribute reduces to + output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter). """ - def shard_parameters(self, module, mesh): - """Wrap selected params as DTensor placeholders. Default: no-op.""" - pass + def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): + self.input_layouts = input_layouts or Shard(-1) + self.output_layouts = output_layouts or Replicate() + self.use_local_output = use_local_output - def transform_inputs_pre_forward(self, module, args, kwargs, mesh): - return args, kwargs + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + if isinstance(module, torch.nn.Embedding): + placement = Shard(0) + else: + # bias is replicated (added after the row-reduce); weight shards on input dim + placement = Replicate() if param == "bias" else Shard(1) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) - def context_around_forward(self, module): - return contextlib.nullcontext() + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + # Embedding runtime sharding needs a replicated input; Linear needs Shard(-1). + desired = Replicate() if isinstance(module, torch.nn.Embedding) else Shard(-1) + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if x.placements != (desired,): + x = x.redistribute(placements=[desired], async_op=True) + return (x,) + args[1:], kwargs def transform_output_post_forward(self, module, output, mesh): - return output + if not isinstance(output, DTensor): + return output + if output.placements != (self.output_layouts,): + output = output.redistribute(placements=[self.output_layouts], async_op=True) + return output.to_local() if self.use_local_output else output - def _install_forward(self, module, mesh): - """Install pre / around / post transforms by replacing module.forward.""" - original_forward = module.forward + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " + f"output_layouts={self.output_layouts}, use_local_output={self.use_local_output})" + ) - def tp_forward(*args, **kwargs): - args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh) - with self.context_around_forward(module): - output = original_forward(*args, **kwargs) - return self.transform_output_post_forward(module, output, mesh) - module.forward = tp_forward - return module +class SequenceParallel(TensorParallelLayer): + def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True): + self.sequence_dim = sequence_dim + self.use_local_output = use_local_output + + def install_forward(self, module, mesh): + # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe). + for p_name, p in list(module.named_parameters(recurse=False)): + module.register_parameter( + p_name, torch.nn.Parameter(DTensor.from_local(p, mesh, [Replicate()], run_check=False)) + ) + return super().install_forward(module, mesh) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + seq = Shard(self.sequence_dim) + x = args[0] + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [seq], run_check=False) + elif x.placements != (seq,): + x = x.redistribute(placements=[seq], async_op=True) + return (x,) + args[1:], kwargs - def _apply(self, module, mesh): - self.shard_parameters(module, mesh) - return self._install_forward(module, mesh) + def transform_output_post_forward(self, module, output, mesh): + if isinstance(output, DTensor): + return output.to_local() if self.use_local_output else output + return output + def __repr__(self) -> str: + return f"{self.__class__.__name__}(sequence_dim={self.sequence_dim}, use_local_output={self.use_local_output})" -class PrepareModuleInputOutput(CustomParallelStyle): - """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)). - Used for MoE blocks with SP: the input sequence is gathered before routing, - and the output (after expert allreduce) is split back to match the residual. - Forward output split is a local op (no comm). Backward creates the all-gather. - """ +class PrepareModuleInputOutput(TensorParallelLayer): + """Allgather input (Shard(1) → Replicate) + local split output (Replicate → Shard(1)).""" def __init__(self, use_local_output=True): - super().__init__() self.use_local_output = use_local_output def transform_inputs_pre_forward(self, module, args, kwargs, mesh): @@ -195,73 +295,67 @@ def transform_output_post_forward(self, module, output, mesh): return output.redistribute(placements=[Shard(1)]).to_local() +class PrepareModuleInput(TensorParallelLayer): + """Allgather a module input (default input_layout → desired_layout, then to-local).""" + + def __init__(self, *, input_kwarg=None, input_layout=None, desired_layout=None, use_local_output=True): + self.input_kwarg = input_kwarg + self.input_layout = input_layout or Shard(1) + self.desired_layout = desired_layout or Replicate() + self.use_local_output = use_local_output + + def _prepare(self, x, mesh): + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [self.input_layout], run_check=False) + if x.placements != (self.desired_layout,): + x = x.redistribute(placements=[self.desired_layout]) + return x.to_local() if self.use_local_output else x + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + if self.input_kwarg is not None: + if kwargs.get(self.input_kwarg) is not None: + kwargs = {**kwargs, self.input_kwarg: self._prepare(kwargs[self.input_kwarg], mesh)} + return args, kwargs + return (self._prepare(args[0], mesh),) + args[1:], kwargs + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(input_kwarg={self.input_kwarg!r}, input_layout={self.input_layout}, " + f"desired_layout={self.desired_layout}, use_local_output={self.use_local_output})" + ) + + +# ============================================================================= +# MoE / packed-linear local-param swap (grouped_mm needs plain tensors) +# ============================================================================= + + def _accumulate_local_param_grad(original_param: DTensor, local_grad: torch.Tensor) -> torch.Tensor: """Stitch a local grad back onto the original DTensor parameter. - During forward we replace the DTensor param with a detached plain-tensor - leaf (see ``_swap_dtensor_params_for_local``) because ``grouped_mm`` / fused - ops don't accept DTensor inputs. That swap breaks the autograd link between - the local leaf's grad and the DTensor param's ``.grad``, so this tensor hook - runs on the leaf and copies/accumulates the grad onto the original DTensor. - - NOTE: An autograd-aware ``param.to_local()`` swap would let backward stitch - the grad automatically, but DTensor's backward path then redistributes the - resulting grad — and that redistribute does not currently support - ``_StridedShard`` placements (used by ``MoEExpertsParallel`` / ``PackedColwiseParallel``). + Used when the forward swap detaches the local leaf (``_StridedShard`` params) because + DTensor backward redistribute does not support that placement as a source. """ tensor_meta = original_param._spec.tensor_meta detached_grad = local_grad.detach() - grad_dtensor = DTensor.from_local( - detached_grad, - original_param.device_mesh, - original_param.placements, - run_check=False, - shape=tensor_meta.shape, - stride=tensor_meta.stride, - ) with torch.no_grad(): - existing_grad = original_param.grad - if existing_grad is None: - original_param.grad = grad_dtensor - elif isinstance(existing_grad, DTensor): - existing_grad._local_tensor.add_(detached_grad) + if original_param.grad is None: + original_param.grad = DTensor.from_local( + detached_grad, + original_param.device_mesh, + original_param.placements, + run_check=False, + shape=tensor_meta.shape, + stride=tensor_meta.stride, + ) + elif isinstance(original_param.grad, DTensor): + original_param.grad._local_tensor.add_(detached_grad) else: - existing_grad.add_(detached_grad) + original_param.grad.add_(detached_grad) return local_grad -@contextlib.contextmanager -def _swap_dtensor_params_for_local(module): - """Temporarily replace DTensor params with local-shard ``Parameter``s for forward. - - ``grouped_mm`` / fused kernels don't accept DTensor inputs, so each DTensor - param is swapped for a detached local ``Parameter``. A tensor hook - (``_accumulate_local_param_grad``) on the local leaf copies the backward - grad back onto the original DTensor. - - The original DTensor params are restored on exit (even on exception) so - save_pretrained / state-dict still see sharded params. - """ - shadows = {} - for name, param in list(module.named_parameters(recurse=False)): - if not isinstance(param, DTensor): - continue - shadows[name] = param - local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) - if param.requires_grad: - local.register_hook(lambda g, p=param: _accumulate_local_param_grad(p, g)) - module._parameters.pop(name) - setattr(module, name, local) - try: - yield - finally: - for name, param in shadows.items(): - if hasattr(module, name): - delattr(module, name) - module.register_parameter(name, param) - - -class PackedColwiseParallel(CustomParallelStyle): +class PackedColwiseParallel(TensorParallelLayer): """Column-wise parallel style for fused linear weights packed along the output dimension.""" def __init__( @@ -271,11 +365,23 @@ def __init__( use_local_output: bool = True, split_factor: int = 2, ): - super().__init__() self.input_layouts = (input_layouts or Replicate(),) self.use_local_output = use_local_output self.split_factor = split_factor + def shard_param(self, module, param, mesh): + if not isinstance(module, torch.nn.Linear): + raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") + meta = module._parameters.get(param) + if meta is None: + return + # Wrap as a DTensor placeholder. Runs on meta — distribute_tensor builds metadata only. + placement = _StridedShard(dim=0, split_factor=self.split_factor) + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): input_tensor = args[0] if not isinstance(input_tensor, DTensor): @@ -285,8 +391,28 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh): input_tensor = input_tensor.to_local() return (input_tensor,) + args[1:], kwargs + @contextlib.contextmanager def context_around_forward(self, module): - return _swap_dtensor_params_for_local(module) + # grouped_mm etc needs plain tensors, so swap the params + to_swap_params = [ + (name, param) for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor) + ] + for name, param in to_swap_params: + del module._parameters[name] + if _find_strided_shard_placement_from_fused_params(param.placements) is not None: + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(functools.partial(_accumulate_local_param_grad, param)) + else: + local = param.to_local() + setattr(module, name, local) + try: + yield + finally: + for name, param in to_swap_params: + # restore the original DTensor params + delattr(module, name) + module.register_parameter(name, param) def transform_output_post_forward(self, module, output, mesh): if output is None or self.use_local_output: @@ -295,21 +421,6 @@ def transform_output_post_forward(self, module, output, mesh): output, mesh, (_StridedShard(dim=-1, split_factor=self.split_factor),), run_check=False ) - def shard_parameters(self, module, mesh): - if not isinstance(module, torch.nn.Linear): - raise NotImplementedError("PackedColwiseParallel currently only supports nn.Linear!") - # Wrap weight + bias as DTensor placeholders. Runs on meta — - # distribute_tensor builds metadata only, no collective. - placement = _StridedShard(dim=0, split_factor=self.split_factor) - for name in ("weight", "bias"): - meta = module._parameters.get(name) - if meta is None: - continue - module._parameters[name] = torch.nn.Parameter( - distribute_tensor(meta, mesh, [placement], src_data_rank=None), - requires_grad=meta.requires_grad, - ) - def __repr__(self) -> str: return ( f"{self.__class__.__name__}(input_layouts={self.input_layouts}, " @@ -317,6 +428,37 @@ def __repr__(self) -> str: ) +class MoEParamShard(TensorParallelLayer): + """Param-only TP/EP style for MoE expert weights: shard the named 3D params as DTensor + placeholders; the matching forward comm is a separate moe_experts_allreduce entry. + + shards_expert_dim (EP grouped_gemm) shards dim 0 and updates module.num_experts to the + per-rank local count, so the experts forward and ep_router sentinel agree. + """ + + def __init__(self, placement, *, shards_expert_dim: bool = False): + self.placement = placement + self.shards_expert_dim = shards_expert_dim + + def shard_param(self, module, param, mesh): + # Wrap one expert param as a DTensor placeholder. Runs on meta — + # distribute_tensor builds metadata only, no collective. + meta = module._parameters.get(param) + if meta is None: + return + if self.shards_expert_dim: + # dim 0 is the expert dimension; record the per-rank local count so the + # experts forward and the EP router agree on local expert ids/sentinel. + module.num_experts = meta.shape[0] // mesh.size() + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [self.placement], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(placement={self.placement}, shards_expert_dim={self.shards_expert_dim})" + + if is_torch_available() and is_torch_greater_or_equal("2.5"): class _AllReduceBackward(torch.autograd.Function): @@ -325,7 +467,7 @@ class _AllReduceBackward(torch.autograd.Function): Used for MoE routing weights: the forward value is replicated (same on all ranks), but the backward gradient is partial (each rank has 1/tp_size from its expert shard). We need to sum the partial gradients without dividing by - world_size, which is what DTensor's ``Replicate`` backward does incorrectly. + world_size, which is what DTensor's Replicate backward does incorrectly. """ @staticmethod @@ -339,59 +481,68 @@ def backward(ctx, grad): return grad, None -class MoEExpertsParallel(CustomParallelStyle): - """Tensor-parallel style for MoE expert modules. - - Shards expert weights as DTensors, then wraps the module's forward so - that grouped_mm (which needs plain tensors) works transparently. - - Lifecycle phases: - 1. shard_parameters — wrap each expert weight named in shard_plan as a DTensor - placeholder with the declared placement. - 2. transform_inputs_pre_forward — localize hidden_states (Replicate→local, - gives us an all-reduce on the backward gradient for free), then fix - routing-weight gradients (their backward is partial; use allreduce-sum, - not divide-by-world-size). - 3. context_around_forward — swap DTensor params for local leaves so - grouped_mm sees plain tensors; restored on exit so save_pretrained - still sees DTensors. - 4. transform_output_post_forward — under TP-only each rank's output is - partial (only its expert shard contributed); reduce/redistribute to - output_layouts. - """ - - def __init__(self, output_layouts=None, shard_plan=None): - super().__init__() +class MoEExpertsParallel(TensorParallelLayer): + def __init__(self, output_layouts=None): self.output_layouts = output_layouts or Replicate() - self._moe_shard_plan = shard_plan or {} - - def shard_parameters(self, module, mesh): - # Wrap each expert weight as a DTensor placeholder. Runs on meta — - # distribute_tensor builds metadata only, no collective. - for name, placement in self._moe_shard_plan.items(): - meta = module._parameters.get(name) - if meta is None: - continue - module._parameters[name] = torch.nn.Parameter( - distribute_tensor(meta, mesh, [placement], src_data_rank=None), - requires_grad=meta.requires_grad, - ) - def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False): hidden_states, top_k_index, top_k_weights = args + tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") if not isinstance(hidden_states, DTensor): hidden_states = DTensor.from_local(hidden_states, mesh, [Replicate()], run_check=False) hidden_states = hidden_states.to_local() + hidden_states = _AllReduceBackward.apply(hidden_states, tp_group) if isinstance(top_k_weights, DTensor): top_k_weights = top_k_weights.to_local() - tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") - top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) + # Under TP the router runs replicated, so its routing-weight gradient is partial + # across ranks and needs an all-reduce-sum in backward. Under EP the router output + # is already sliced per-rank by `ep_router` (non-local scores zeroed), so each + # rank's routing weights are independent — skip the all-reduce to avoid + # double-counting the gradient. + if not is_expert_parallel: + top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group) return (hidden_states, top_k_index, top_k_weights), kwargs + def install_forward(self, module, mesh, *, is_expert_parallel=False): + """Install pre / around / post transforms; ``is_expert_parallel`` is baked into the closure.""" + original_forward = module.forward + + def tp_forward(*args, **kwargs): + args, kwargs = self.transform_inputs_pre_forward( + module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel + ) + with self.context_around_forward(module): + output = original_forward(*args, **kwargs) + return self.transform_output_post_forward(module, output, mesh) + + module.forward = tp_forward + return module + + @contextlib.contextmanager def context_around_forward(self, module): - return _swap_dtensor_params_for_local(module) + # grouped_mm experts forward needs plain tensors, so swap the params + to_swap_params = [ + (name, param) for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor) + ] + for name, param in to_swap_params: + del module._parameters[name] + # Happens only when train in TP only (experts dont use grouped_gemm but rely on moe_tp_gate_up_colwise) + if _find_strided_shard_placement_from_fused_params(param.placements) is not None: + local = torch.nn.Parameter(param._local_tensor.detach(), requires_grad=param.requires_grad) + if param.requires_grad: + local.register_hook(functools.partial(_accumulate_local_param_grad, param)) + else: + local = param.to_local() + setattr(module, name, local) + try: + yield + finally: + # restore the original DTensor params + for name, param in to_swap_params: + delattr(module, name) + module.register_parameter(name, param) def transform_output_post_forward(self, module, output, mesh): if output is None: @@ -414,19 +565,85 @@ def transform_output_post_forward(self, module, output, mesh): return output.to_local() +class MoeIdentityParallel(TensorParallelLayer): + """ + TP class for zero/identity experts in MoE layers. + + Under TP, the parent `MoEExpertsParallel` all-reduces the expert module + output by summation. Identity experts produce the same output on every + rank, so the sum gives `world_size * output`. This class divides the input + by `world_size` to compensate. + """ + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + input_tensor = args[0] + return (input_tensor / mesh.size(), *args[1:]), kwargs + + +class EpRouterParallel(TensorParallelLayer): + """Expert-parallel router: forward-only slicing of router outputs to local experts. + + Ported from the original RouterParallel (#39501). The gate runs replicated on + every rank and emits global (router_logits, router_scores, router_indices). Under + EP each rank owns num_experts // ep_size experts, so this post-forward hook zeroes + the scores of non-local experts and remaps the surviving indices into the local range, + using num_local_experts as the sentinel for dropped (non-local) slots. The + downstream experts forward masks that sentinel and moe_experts_allreduce sums the + partial per-rank results. + + No parameter sharding — the gate weight stays replicated. + + Example: 128 experts, EP=8 → num_local_experts=16. Rank 0 (owns experts 0-15) keeps + only indices in [0, 16), remaps them via fmod, and sets everything else to the + sentinel 16; its scores are zeroed everywhere except the surviving local slots. + """ + + def transform_output_post_forward(self, module, output, mesh): + ep_rank, ep_size = mesh.get_local_rank(), mesh.size() + num_experts = getattr(module, "num_experts", None) + if num_experts is None: + num_experts = getattr(getattr(module, "config", None), "num_experts", None) + if num_experts is None: + raise AttributeError( + f"Router module {type(module).__name__} is missing `num_experts` and `config.num_experts`" + ) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts must be divisible by ep_size: {num_experts} % {ep_size} != 0") + num_local_experts = num_experts // ep_size + + router_logits, router_scores, router_indices = output + non_local_mask = (router_indices // num_local_experts) != ep_rank + router_scores = router_scores.masked_fill(non_local_mask, 0.0) + router_indices = router_indices.masked_fill(non_local_mask, -1) + # `-1 % 1 == 0`, so fmod only remaps correctly when there is more than one local expert. + if num_local_experts > 1: + router_indices = torch.fmod(router_indices, num_local_experts) + else: + router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1) + router_indices = router_indices.masked_fill(router_indices == -1, num_local_experts) + return router_logits, router_scores, router_indices + + class ParallelInterface(GeneralInterface): """Registry of named TP styles. Configs and modeling files reference these by string name. - Adding a new entry here is the supported way to introduce a new TP style. - Users can also override or extend at runtime via ``ALL_PARALLEL_STYLES["my_style"] = ...``. + Styles are split into the three parallelism groups they belong to (TENSOR_/SEQUENCE_/ + EXPERT_PARALLEL_STYLES); _global_mapping is the union of all three. Keeping them separate + lets other code reason about a group as a whole — e.g. SEQUENCE_PARALLEL_STYLES is reused to + validate that a user-provided tp_plan actually shards the sequence dimension when + enable_sequence_parallel=True. + + Adding a new entry under one of the three groups is the supported way to introduce a new TP + style. Users can also override or extend at runtime via ALL_PARALLEL_STYLES["my_style"] = .... - Naming convention: ``{kind}[_{comm}][_{extra}]``. The ``_{comm}`` suffix is dropped only when - comm is ``"none"`` (no collective). All entries are eager instances; the dict literal lives - behind a torch-availability guard so this module remains importable without torch. + Naming convention: {kind}[_{comm}][_{extra}]. The _{comm} suffix is dropped only when + comm is "none" (no collective). All entries are eager instances, so the dicts live behind a + torch-availability guard and are empty when torch is unavailable (this module stays + importable; styles are only ever looked up inside apply_tensor_parallel, which needs torch). """ - _global_mapping = ( - { + if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available: + TENSOR_PARALLEL_STYLES = { # Column-parallel "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), "colwise_allgather": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), @@ -434,80 +651,133 @@ class ParallelInterface(GeneralInterface): input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False ), "packed_colwise": PackedColwiseParallel(input_layouts=Replicate()), + # MoE intra-expert weight sharding — param-level (no forward hook). The matching forward + # comm is declared separately via "moe_experts_allreduce" on the experts module. gate_up + # is packed (gate||up) along dim -2 (Qwen3/Mixtral: [E, 2*inter, hidden]). + "moe_tp_gate_up_colwise": MoEParamShard(_StridedShard(dim=-2, split_factor=2)), # packed gate/up + "moe_tp_down_rowwise": MoEParamShard(Shard(-1)), # down_proj input dim # Row-parallel "rowwise_allreduce": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()), - "rowwise_reduce_scatter": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Shard(1)), # Vocab / embedding (rowwise sharding on vocab dim) "vocab_allreduce": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + } + SEQUENCE_PARALLEL_STYLES = { + "rowwise_reduce_scatter": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Shard(1)), "vocab_reduce_scatter": RowwiseParallel(input_layouts=Replicate(), output_layouts=Shard(1)), - # Activation / norm (sequence-parallel passthrough) - # use_local_output=True: torch defaults to False here, but downstream modeling - # code expects plain tensors, not DTensors. + # Activation / norm (sequence-parallel passthrough). use_local_output=True: torch defaults + # to False here, but downstream modeling code expects plain tensors, not DTensors. "activation": SequenceParallel(use_local_output=True), "activation_seq_dim_2": SequenceParallel(sequence_dim=2, use_local_output=True), - # Module-level prepare-input. Same use_local_output=True override as above — - # torch's default is False, our modeling code expects plain tensors downstream. - "module_allgather": PrepareModuleInput( - input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),), use_local_output=True - ), + # Module-level prepare-input (allgather seq-sharded input → replicate, to local). + # use_local_output=True: our modeling code expects plain tensors downstream. + "module_allgather": PrepareModuleInput(input_layout=Shard(1), desired_layout=Replicate()), "module_allgather_hidden_states": PrepareModuleInput( - input_kwarg_layouts={"hidden_states": Shard(1)}, - desired_input_kwarg_layouts={"hidden_states": Replicate()}, - use_local_output=True, + input_kwarg="hidden_states", input_layout=Shard(1), desired_layout=Replicate() ), "module_allgather_split": PrepareModuleInputOutput(), - # MoE — canonical shard_plan baked in (only variant in use across configs). - # gate_up_proj is packed (gate||up along output dim) so we use _StridedShard - # to interleave; down_proj is plain rowwise on its input dim. - "moe_experts_allreduce": MoEExpertsParallel( - output_layouts=Replicate(), - shard_plan={ - "gate_up_proj": _StridedShard(dim=-2, split_factor=2), - "down_proj": Shard(-1), - }, - ), } - if is_torch_available() and is_torch_greater_or_equal("2.5") and _torch_distributed_available - else {} - ) + EXPERT_PARALLEL_STYLES = { + "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), # expert dim + # Experts' forward comm (no weight sharding — that is declared per-param via "grouped_gemm" + # / "moe_tp_*"). Installs the forward hook that all-reduces partial per-rank expert outputs. + # Reused by dense TP-MoE as well as EP. + "moe_experts_allreduce": MoEExpertsParallel(output_layouts=Replicate()), + "moe_identity_expert": MoeIdentityParallel(), + # EP router — forward-only slicing of router outputs to local experts. + "ep_router": EpRouterParallel(), + } + else: + TENSOR_PARALLEL_STYLES = SEQUENCE_PARALLEL_STYLES = EXPERT_PARALLEL_STYLES = {} + + _global_mapping = {**TENSOR_PARALLEL_STYLES, **SEQUENCE_PARALLEL_STYLES, **EXPERT_PARALLEL_STYLES} ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface() -def apply_tensor_parallel(model, tp_mesh, tp_plan): - """Apply tensor parallelism by calling each style's ``_apply`` on the - matching submodules. +def select_parallel_plan(model) -> dict[str, str]: + """ + Select the parallel plan to apply from explicit combo plans on the model. + """ + distributed_config = model.config.distributed_config + user_tp_plan = distributed_config.tp_plan + if user_tp_plan is not None: + return dict(user_tp_plan) + + sp = distributed_config.enable_sequence_parallel + ep = distributed_config.enable_expert_parallel + if sp and ep: + plan_attr = "_sp_ep_plan" + elif sp: + plan_attr = "_sp_plan" + elif ep: + plan_attr = "_tp_ep_plan" + else: + plan_attr = "_tp_plan" + config_attr = f"base_model{plan_attr}" + + plan = getattr(model, plan_attr, None) or {} + if not plan: + raise ValueError( + f"Model {model.config.model_type!r} has no {config_attr} (model.{plan_attr} is empty) but " + f"DistributedConfig requested enable_sequence_parallel={sp}, enable_expert_parallel={ep}. " + f"Add {config_attr} to the model config, set DistributedConfig.tp_plan explicitly, " + f"or change the flags." + ) + return dict(plan) + + +def apply_tensor_parallel(model, tp_mesh): + """Apply tensor parallelism by looking each style up by string name. - Walks ``model.named_modules()``, resolves each name against the wildcard - ``tp_plan`` from the model config, and applies the corresponding style - from ``ALL_PARALLEL_STYLES`` (looked up by string name) directly. + A single named_modules() walk shards each module's direct params (param-keyed plan + entries → shard_param) and installs its forward comm (module-keyed entries → + install_forward). Plus the cross-cutting setup: tie-weights handling, SP position_ids + injection, and loss_parallel activation. """ - distributed_config = getattr(model.config, "distributed_config", None) - sp_requested = getattr(distributed_config, "enable_sequence_parallel", False) - sp_supported = getattr(model.config, "base_model_sp_plan", None) is not None - enable_sp = sp_requested and sp_supported + enable_sp = bool(getattr(model.config.distributed_config, "enable_sequence_parallel", False)) + enable_ep = bool(getattr(model.config.distributed_config, "enable_expert_parallel", False)) + + # case when enable_sequence_parallel=True and tp_plan!=None (but tp_plan has no sequence parallel styles keys) raise an error + user_tp_plan = model.config.distributed_config.tp_plan + if ( + enable_sp + and user_tp_plan is not None # user provided tp_plan manually + and not any(style in ParallelInterface.SEQUENCE_PARALLEL_STYLES for style in user_tp_plan.values()) + ): + raise ValueError( + "enable_sequence_parallel=True but the provided tp_plan shards no sequence dimension " + f"(none of {sorted(ParallelInterface.SEQUENCE_PARALLEL_STYLES)}). Provide a sequence-parallel plan " + "(e.g. 'vocab_reduce_scatter' on the embedding) or set enable_sequence_parallel=False. " + f"tp_plan: {user_tp_plan}" + ) - if tp_plan is None: - if enable_sp: - tp_plan = dict(model._sp_plan or {}) - else: - tp_plan = dict(model._tp_plan or {}) + model.tp_plan = dict(select_parallel_plan(model)) + logger.info(f"TP plan has been resolved: {model.tp_plan}") # tie_weights() replaces lm_head.weight with embed_tokens.weight after TP is applied. # If embed_tokens isn't in the plan, sharding lm_head as a DTensor causes tie to # clobber it with a plain tensor (and forward then mixes DTensor/Tensor). Skip # lm_head TP in that case so both ends stay plain and the tie is a real alias. if getattr(model.config, "tie_word_embeddings", False): - tied_source_in_plan = any(k.endswith("embed_tokens") for k in tp_plan) + tied_source_in_plan = any(k.endswith("embed_tokens") for k in model.tp_plan) if not tied_source_in_plan: - tp_plan.pop("lm_head", None) - - for name, submodule in model.named_modules(): - style_value = _get_parameter_tp_plan(parameter_name=name, tp_plan=tp_plan, is_weight=False) - if style_value is None: - continue - ALL_PARALLEL_STYLES[style_value]._apply(submodule, tp_mesh) + model.tp_plan.pop("lm_head", None) + + for name, module in model.named_modules(): + # Shard each parameter directly on the module using the specified style. + for p_name, _ in list(module.named_parameters(recurse=False)): + full = f"{name}.{p_name}" if name else p_name + style_name = _get_parameter_tp_plan(parameter_name=full, tp_plan=model.tp_plan, is_weight=True) + if style_name is not None and style_name in ALL_PARALLEL_STYLES: + ALL_PARALLEL_STYLES[style_name].shard_param(module, p_name, tp_mesh) + # Install forward hooks for modules as needed by the plan. + style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=model.tp_plan, is_weight=False) + if style_name is not None and style_name in ALL_PARALLEL_STYLES: + if style_name == "moe_experts_allreduce": + ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh, is_expert_parallel=enable_ep) + else: + ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh) # Under SP, inputs_embeds is sequence-sharded after embed_tokens, so # auto-generated position_ids would use the wrong (local) seq_len. @@ -531,7 +801,7 @@ def _inject_sp_metadata(mod, args, kwargs): # loss_parallel patches F.cross_entropy to work with Shard(-1) logits. # It must be active during both forward and backward, so we enable it # once rather than as a context manager. - has_loss_parallel = any(v == "colwise_loss_parallel" for v in tp_plan.values()) + has_loss_parallel = any(v == "colwise_loss_parallel" for v in model.tp_plan.values()) if has_loss_parallel: from torch.distributed.tensor.parallel import loss_parallel diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 855d5f8175dc..4146058a7e84 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -146,7 +146,7 @@ def distribute_model(model, distributed_config: DistributedConfig, device_mesh) mesh_dim_names = device_mesh.mesh_dim_names or () if "tp" in mesh_dim_names: tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh - model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan) + model = apply_tensor_parallel(model, tp_mesh) if "fsdp" in mesh_dim_names: fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh model = apply_fully_shard_data_parallel(model, fsdp_mesh) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 76bc90d20cee..0bd4786e28e4 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -443,6 +443,12 @@ def grouped_mm_experts_forward( proj_out = _grouped_linear( selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating + # Mid-mask (bwd path). The up grouped_mm leaves sentinel rows of `proj_out` uninit. This + # matters in BACKWARD: the masked_fill_ backward zeros the gradient at sentinel rows (select, + # so even NaN → 0) before the non-linearity bwd's garbage can reach grad(gate_up_proj[_bias]). + # Forward needs no masking here (the down grouped_mm skips these rows) — the value-zeroing is + # just a harmless side effect. + proj_out.masked_fill_(sentinel_mask, 0.0) # Apply gating or activation if self.has_gate: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index bd8869e9f3f9..d86cf4811714 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1410,10 +1410,13 @@ def post_init(self): # before the instance attribute shadows them — task-head plans live on the head class. cls_tp_plan = getattr(self, "_tp_plan", None) or {} cls_sp_plan = getattr(self, "_sp_plan", None) or {} + cls_tp_ep_plan = getattr(self, "_tp_ep_plan", None) or {} + cls_sp_ep_plan = getattr(self, "_sp_ep_plan", None) or {} cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} self._tp_plan = dict(cls_tp_plan) self._sp_plan = dict(cls_sp_plan) - self._ep_plan = {} + self._tp_ep_plan = dict(cls_tp_ep_plan) + self._sp_ep_plan = dict(cls_sp_ep_plan) self._pp_plan = {} self._fsdp_plan = dict(cls_fsdp_plan) # If current model is a base model, attach `base_model_*_plan` from config @@ -1421,7 +1424,12 @@ def post_init(self): self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} self._sp_plan = self.config.base_model_sp_plan.copy() if self.config.base_model_sp_plan is not None else {} - self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {} + self._tp_ep_plan = ( + self.config.base_model_tp_ep_plan.copy() if self.config.base_model_tp_ep_plan is not None else {} + ) + self._sp_ep_plan = ( + self.config.base_model_sp_ep_plan.copy() if self.config.base_model_sp_ep_plan is not None else {} + ) self._fsdp_plan = ( self.config.base_model_fsdp_plan.copy() if self.config.base_model_fsdp_plan is not None else {} ) @@ -1443,12 +1451,14 @@ def post_init(self): # This works because the way the `__init__` and `post_init` are called on all submodules is depth-first in the graph for name, module in self.named_children(): # Parallel plans - if plan := getattr(module, "_ep_plan", None): - self._ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_tp_plan", None): self._tp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_sp_plan", None): self._sp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_tp_ep_plan", None): + self._tp_ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_sp_ep_plan", None): + self._sp_ep_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_pp_plan", None): self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_fsdp_plan", None): @@ -1482,11 +1492,6 @@ def post_init(self): @property def tp_plan(self) -> dict[str, str]: - """ - The full tp plan for the model's modules - """ - if hasattr(self.config, "distributed_config") and self.config.distributed_config.enable_expert_parallel: - return self._ep_plan return self._tp_plan @property @@ -4395,10 +4400,6 @@ def from_pretrained( # Expand device_map if it was passed as a `str`, i.e. `device_map="auto"` device_map = _get_device_map(model, device_map, max_memory, hf_quantizer) - # Finalize model weight initialization - active_tp_plan = ( - getattr(model, "_tp_plan", None) if getattr(distributed_config, "tp_size", None) is not None else None - ) load_config = LoadStateDictConfig( pretrained_model_name_or_path=pretrained_model_name_or_path, ignore_mismatched_sizes=ignore_mismatched_sizes, @@ -4410,7 +4411,7 @@ def from_pretrained( dtype_plan=dtype_plan, hf_quantizer=hf_quantizer, device_mesh=device_mesh, - tp_plan=active_tp_plan, + tp_plan=model.tp_plan, fsdp_plan=model.fsdp_plan, weights_only=weights_only, weight_mapping=weight_conversions, diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 1fb30e564e1c..cf4a595bdcda 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -617,6 +617,8 @@ class AfmoeForCausalLM(AfmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index 0875127a3215..b055b222e38d 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -398,6 +398,8 @@ class AfmoeForCausalLM(LlamaForCausalLM, AfmoePreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/apertus/modeling_apertus.py b/src/transformers/models/apertus/modeling_apertus.py index 8c5757377b03..85bb6b27fd23 100644 --- a/src/transformers/models/apertus/modeling_apertus.py +++ b/src/transformers/models/apertus/modeling_apertus.py @@ -424,6 +424,8 @@ class ApertusForCausalLM(ApertusPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/arcee/modeling_arcee.py b/src/transformers/models/arcee/modeling_arcee.py index 05ebab40a879..5aaee35657a2 100644 --- a/src/transformers/models/arcee/modeling_arcee.py +++ b/src/transformers/models/arcee/modeling_arcee.py @@ -426,6 +426,8 @@ class ArceeForCausalLM(ArceePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index c17f1541484f..f23cf5d27fca 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -760,6 +760,8 @@ class AriaTextForCausalLM(AriaTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/bamba/modeling_bamba.py b/src/transformers/models/bamba/modeling_bamba.py index d59423cf38a4..77a49ebe7d00 100644 --- a/src/transformers/models/bamba/modeling_bamba.py +++ b/src/transformers/models/bamba/modeling_bamba.py @@ -1072,6 +1072,8 @@ class BambaForCausalLM(BambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/bitnet/modeling_bitnet.py b/src/transformers/models/bitnet/modeling_bitnet.py index 4994407d74c1..b51bd8fdabf1 100644 --- a/src/transformers/models/bitnet/modeling_bitnet.py +++ b/src/transformers/models/bitnet/modeling_bitnet.py @@ -424,6 +424,8 @@ class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None _sp_plan = None + _tp_ep_plan = None + _sp_ep_plan = None _pp_plan = None _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/bitnet/modular_bitnet.py b/src/transformers/models/bitnet/modular_bitnet.py index c8376f131859..9bb7991e6d72 100644 --- a/src/transformers/models/bitnet/modular_bitnet.py +++ b/src/transformers/models/bitnet/modular_bitnet.py @@ -111,6 +111,8 @@ class BitNetForCausalLM(LlamaForCausalLM): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = None _sp_plan = None + _tp_ep_plan = None + _sp_ep_plan = None _pp_plan = None def forward( diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index 610fec7f0e84..6d4a6de09588 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -454,6 +454,8 @@ class CohereForCausalLM(CoherePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/cohere2/modeling_cohere2.py b/src/transformers/models/cohere2/modeling_cohere2.py index 051867c8f8d9..e78e1adaae6b 100644 --- a/src/transformers/models/cohere2/modeling_cohere2.py +++ b/src/transformers/models/cohere2/modeling_cohere2.py @@ -432,6 +432,8 @@ class Cohere2ForCausalLM(Cohere2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py index f48a04a5939f..76050d934197 100644 --- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py @@ -76,6 +76,8 @@ class Cohere2MoeConfig(PreTrainedConfig): "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", # Depending on layers, `mlp` can be a MoE or an MLP layer + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", diff --git a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py index 94b1de772fc2..20b4abef6c37 100644 --- a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py +++ b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py @@ -582,6 +582,8 @@ class Cohere2MoeForCausalLM(Cohere2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/csm/modeling_csm.py b/src/transformers/models/csm/modeling_csm.py index 99f643e93b44..97f5bf6425aa 100644 --- a/src/transformers/models/csm/modeling_csm.py +++ b/src/transformers/models/csm/modeling_csm.py @@ -550,6 +550,8 @@ class CsmDepthDecoderForCausalLM(CsmPreTrainedModel, GenerationMixin): _tied_weights_keys = None _tp_plan = None _sp_plan = None + _tp_ep_plan = None + _sp_ep_plan = None _pp_plan = None _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/csm/modular_csm.py b/src/transformers/models/csm/modular_csm.py index f1e95384ab52..4fa7c150285b 100644 --- a/src/transformers/models/csm/modular_csm.py +++ b/src/transformers/models/csm/modular_csm.py @@ -275,6 +275,8 @@ class CsmDepthDecoderForCausalLM(LlamaForCausalLM, GenerationMixin): _tied_weights_keys = None _tp_plan = None _sp_plan = None + _tp_ep_plan = None + _sp_ep_plan = None _pp_plan = None def __init__(self, config): diff --git a/src/transformers/models/cwm/modeling_cwm.py b/src/transformers/models/cwm/modeling_cwm.py index b20fd9b7436f..fd4e9176aef8 100644 --- a/src/transformers/models/cwm/modeling_cwm.py +++ b/src/transformers/models/cwm/modeling_cwm.py @@ -428,6 +428,8 @@ class CwmForCausalLM(CwmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index a984b4585a55..a425a78ff82d 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -645,6 +645,8 @@ class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/dbrx/modular_dbrx.py b/src/transformers/models/dbrx/modular_dbrx.py index 548af86ef4f6..ab7b3998b085 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -433,6 +433,8 @@ class DbrxForCausalLM(DbrxPreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: DbrxConfig): diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 49b3cb9c6387..f0093fc52a3f 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -56,6 +56,8 @@ class DeepseekV2Config(PreTrainedConfig): "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", @@ -76,6 +78,8 @@ class DeepseekV2Config(PreTrainedConfig): "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", # Shared experts output must stay Replicate to match experts output # (they're summed inside the MoE block, before the outer allgather_split diff --git a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py index 1f9ba33ec6a5..1d4995293a32 100644 --- a/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py @@ -543,6 +543,8 @@ class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index a6c6494c0415..9458009391cb 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py @@ -71,6 +71,8 @@ class DeepseekV2Config(LlamaConfig): "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", @@ -91,6 +93,8 @@ class DeepseekV2Config(LlamaConfig): "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", # Shared experts output must stay Replicate to match experts output # (they're summed inside the MoE block, before the outer allgather_split diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index 32e67ab4d90a..a94ef5660e80 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py @@ -49,6 +49,8 @@ class DeepseekV3Config(PreTrainedConfig): model_type = "deepseek_v3" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index ec44b254bdbf..876a2ccd5a36 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -637,6 +637,8 @@ class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py index baf210824e85..11d12a756eb8 100644 --- a/src/transformers/models/deepseek_v32/modular_deepseek_v32.py +++ b/src/transformers/models/deepseek_v32/modular_deepseek_v32.py @@ -371,7 +371,9 @@ class DeepseekV32Model(DeepseekV3Model): class DeepseekV32ForCausalLM(DeepseekV3ForCausalLM): _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = AttributeError() _sp_plan = AttributeError() + _sp_ep_plan = AttributeError() _fsdp_plan = AttributeError() diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index 5cbea35c1d8f..dcbbb9f0ed47 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -1425,6 +1425,8 @@ class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/src/transformers/models/diffllama/modeling_diffllama.py index b822da50e2af..51f4642d3366 100644 --- a/src/transformers/models/diffllama/modeling_diffllama.py +++ b/src/transformers/models/diffllama/modeling_diffllama.py @@ -662,6 +662,8 @@ class DiffLlamaForCausalLM(DiffLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py b/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py index 07d9d7e2e3bf..69b0d8e606af 100644 --- a/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py +++ b/src/transformers/models/diffusion_gemma/modular_diffusion_gemma.py @@ -113,6 +113,7 @@ class DiffusionGemmaTextConfig(Gemma4TextConfig): "layers.*.experts.down_proj": "grouped_gemm", "layers.*.experts": "moe_tp_experts", } + base_model_tp_ep_plan = AttributeError() base_model_pp_plan = { "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index 1f112465369d..885321f8d197 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -718,6 +718,8 @@ class DogeForCausalLM(DogePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py index 9846f961250f..565fd7b4bd42 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/src/transformers/models/dots1/configuration_dots1.py @@ -52,6 +52,8 @@ class Dots1Config(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index 05b1a6d20b4c..8708a661c2fe 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -571,6 +571,8 @@ class Dots1ForCausalLM(Dots1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index 94bd074871c0..74ae84c5c6b1 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/src/transformers/models/dots1/modular_dots1.py @@ -66,6 +66,8 @@ class Dots1Config(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index fe4202de68c9..4a48b9e60dbd 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -1277,6 +1277,8 @@ class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config: Emu3TextConfig diff --git a/src/transformers/models/ernie4_5/modeling_ernie4_5.py b/src/transformers/models/ernie4_5/modeling_ernie4_5.py index 3f280bc3185d..55691646b590 100644 --- a/src/transformers/models/ernie4_5/modeling_ernie4_5.py +++ b/src/transformers/models/ernie4_5/modeling_ernie4_5.py @@ -423,6 +423,8 @@ class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py index 035daabd966d..5983c9d0fd15 100644 --- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py @@ -67,6 +67,8 @@ class Ernie4_5_MoeConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index 3dc034269a0c..53c04e6bb86b 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -655,6 +655,8 @@ class Ernie4_5_MoeForCausalLM(Ernie4_5_MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py index 34caaa02dbeb..e54aff4408b9 100644 --- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py @@ -87,6 +87,12 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.text_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.text_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.text_moe.experts": "moe_experts_allreduce", + "layers.*.mlp.vision_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.vision_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.vision_moe.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py index 9c05d3943cdf..660d6277eacd 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py @@ -141,6 +141,12 @@ class Ernie4_5_VLMoeTextConfig(Ernie4_5_MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.text_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.text_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.text_moe.experts": "moe_experts_allreduce", + "layers.*.mlp.vision_moe.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.vision_moe.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.vision_moe.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/eurobert/modeling_eurobert.py b/src/transformers/models/eurobert/modeling_eurobert.py index 655d10ab123d..01642bf2be06 100644 --- a/src/transformers/models/eurobert/modeling_eurobert.py +++ b/src/transformers/models/eurobert/modeling_eurobert.py @@ -411,6 +411,8 @@ class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/eurobert/modular_eurobert.py b/src/transformers/models/eurobert/modular_eurobert.py index b40c4d043707..c58246231e40 100644 --- a/src/transformers/models/eurobert/modular_eurobert.py +++ b/src/transformers/models/eurobert/modular_eurobert.py @@ -144,6 +144,8 @@ class EuroBertForMaskedLM(EuroBertPreTrainedModel): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: EuroBertConfig): diff --git a/src/transformers/models/exaone4/modeling_exaone4.py b/src/transformers/models/exaone4/modeling_exaone4.py index d6aa12baedad..4516901f11e3 100644 --- a/src/transformers/models/exaone4/modeling_exaone4.py +++ b/src/transformers/models/exaone4/modeling_exaone4.py @@ -442,6 +442,8 @@ class Exaone4ForCausalLM(Exaone4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 83a873314700..df164851667d 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/src/transformers/models/exaone_moe/configuration_exaone_moe.py @@ -64,12 +64,18 @@ class ExaoneMoeConfig(PreTrainedConfig): model_type = "exaone_moe" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `LlamaModel` + base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/exaone_moe/modeling_exaone_moe.py b/src/transformers/models/exaone_moe/modeling_exaone_moe.py index d2d8543f1417..7f1fcd5fc177 100644 --- a/src/transformers/models/exaone_moe/modeling_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modeling_exaone_moe.py @@ -565,6 +565,8 @@ class ExaoneMoeForCausalLM(ExaoneMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/exaone_moe/modular_exaone_moe.py b/src/transformers/models/exaone_moe/modular_exaone_moe.py index 8d50c3db2295..3e5dfda698ee 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/src/transformers/models/exaone_moe/modular_exaone_moe.py @@ -82,6 +82,22 @@ class ExaoneMoeConfig(Exaone4Config): base_model_sp_plan = None + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + base_model_fsdp_plan = { "embed_tokens": "free_full_weight", "layers.*": "free_full_weight", diff --git a/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/src/transformers/models/falcon_h1/modeling_falcon_h1.py index 2ea3c884437e..823adc24ab84 100644 --- a/src/transformers/models/falcon_h1/modeling_falcon_h1.py +++ b/src/transformers/models/falcon_h1/modeling_falcon_h1.py @@ -1168,6 +1168,8 @@ class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index a89d9f0452e8..a1dc3273f571 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/src/transformers/models/flex_olmo/configuration_flex_olmo.py @@ -54,6 +54,25 @@ class FlexOlmoConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index 517a9b9fcc1c..6f3f8a4af84f 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -605,6 +605,8 @@ class FlexOlmoForCausalLM(FlexOlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index c7b73dea8f15..ea11df43a532 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modular_flex_olmo.py @@ -64,6 +64,25 @@ class FlexOlmoConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k "layers.*.self_attn.v_proj": "colwise_allgather", # we need to replicate here due to the added norm on q and k "layers.*.self_attn.o_proj": "vocab_allreduce", # input is replicated due to the added norm on q and k + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma/modeling_gemma.py b/src/transformers/models/gemma/modeling_gemma.py index f64702b55f19..ad2b3cdcd751 100644 --- a/src/transformers/models/gemma/modeling_gemma.py +++ b/src/transformers/models/gemma/modeling_gemma.py @@ -452,6 +452,8 @@ class GemmaForCausalLM(GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py index 46eb89681619..9a88238fd852 100644 --- a/src/transformers/models/gemma2/modeling_gemma2.py +++ b/src/transformers/models/gemma2/modeling_gemma2.py @@ -478,6 +478,8 @@ class Gemma2ForCausalLM(Gemma2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index a4fefdd2321d..b786a390c7a9 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -587,6 +587,8 @@ class Gemma3ForCausalLM(Gemma3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma3TextConfig diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index 895de17358f4..13482907b5e8 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -1829,6 +1829,8 @@ class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma3nTextConfig diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index c1c8cd1f26d4..a415cc1fc74b 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -136,7 +136,17 @@ class Gemma4TextConfig(PreTrainedConfig): "layers.*.mlp.down_proj": "rowwise_allreduce", } base_model_ep_plan = { - # EP plan for google/gemma-4-26B-A4B-it: do not tp in attention (num_global_key_value_heads=2 too small to partition) + # EP-only: shard routed experts; keep attention and dense MLP replicated. + "layers.*.router": "ep_router", + "layers.*.experts.gate_up_proj": "grouped_gemm", + "layers.*.experts.down_proj": "grouped_gemm", + "layers.*.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/gemma4/modeling_gemma4.py b/src/transformers/models/gemma4/modeling_gemma4.py index 20931ff60c34..350f73651050 100644 --- a/src/transformers/models/gemma4/modeling_gemma4.py +++ b/src/transformers/models/gemma4/modeling_gemma4.py @@ -1819,6 +1819,8 @@ class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config: Gemma4TextConfig diff --git a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py index 61450d24036e..177094e3cd47 100644 --- a/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py +++ b/src/transformers/models/gemma4_assistant/modeling_gemma4_assistant.py @@ -110,6 +110,7 @@ def _init_weights(self, module): class Gemma4AssistantForCausalLM(Gemma4AssistantPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = {"lm_head": "colwise_gather_output"} _fsdp_plan = {"lm_head": "keep_full_weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} diff --git a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py index 1967c45b839b..d3aa2283c59b 100644 --- a/src/transformers/models/gemma4_unified/modular_gemma4_unified.py +++ b/src/transformers/models/gemma4_unified/modular_gemma4_unified.py @@ -456,6 +456,7 @@ class Gemma4UnifiedTextConfig(Gemma4TextConfig): "layers.*.mlp.down_proj": "rowwise", } base_model_ep_plan = None # No MoE + base_model_tp_ep_plan = AttributeError() base_model_fsdp_plan = AttributeError() sliding_window: int = 1024 @@ -759,7 +760,9 @@ def forward( class Gemma4UnifiedForCausalLM(Gemma4ForCausalLM): _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = AttributeError() _sp_plan = AttributeError() + _sp_ep_plan = AttributeError() _fsdp_plan = AttributeError() def forward( diff --git a/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py b/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py index 59f4f3c74e73..c5638a637a6f 100644 --- a/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py +++ b/src/transformers/models/gemma4_unified_assistant/modular_gemma4_unified_assistant.py @@ -16,6 +16,7 @@ class Gemma4UnifiedAssistantForCausalLM(Gemma4AssistantForCausalLM): + _tp_ep_plan = AttributeError() _fsdp_plan = AttributeError() def forward(**super_kwargs): diff --git a/src/transformers/models/glm/modeling_glm.py b/src/transformers/models/glm/modeling_glm.py index 676f5661c1bb..3c17a1159bc6 100644 --- a/src/transformers/models/glm/modeling_glm.py +++ b/src/transformers/models/glm/modeling_glm.py @@ -440,6 +440,8 @@ class GlmForCausalLM(GlmPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/glm4/modeling_glm4.py b/src/transformers/models/glm4/modeling_glm4.py index 2ae82e9da8cf..5f9883211880 100644 --- a/src/transformers/models/glm4/modeling_glm4.py +++ b/src/transformers/models/glm4/modeling_glm4.py @@ -445,6 +445,8 @@ class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py index bd94a492bb66..9856c3ae9dd0 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/src/transformers/models/glm4_moe/configuration_glm4_moe.py @@ -58,6 +58,8 @@ class Glm4MoeConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index 051ccc98ee18..d312f1bd50bc 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -578,6 +578,8 @@ class Glm4MoeForCausalLM(Glm4MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py index 41921eb080e6..1e7266ec206d 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modular_glm4_moe.py @@ -71,6 +71,8 @@ class Glm4MoeConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py index cc0c1fc67ffe..5ab913299ea1 100644 --- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py @@ -55,6 +55,8 @@ class Glm4MoeLiteConfig(PreTrainedConfig): "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", diff --git a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py index c216ee895f3e..3002cbd66fbb 100644 --- a/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py @@ -654,6 +654,8 @@ class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py index ea8e5f60e491..44a369cc8205 100644 --- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py +++ b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py @@ -63,6 +63,8 @@ class Glm4MoeLiteConfig(PreTrainedConfig): "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 73971460b2d4..67399f1868bc 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -57,6 +57,12 @@ class Glm4vMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 289dac43e160..a6515ccd88c8 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -89,6 +89,12 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 0cac376ea2fb..0b2806cd0c65 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -66,6 +66,8 @@ class GlmMoeDsaConfig(PreTrainedConfig, RotaryEmbeddingConfigMixin): "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 46845af56d5f..a618a6d4b6f9 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -82,6 +82,8 @@ class GlmMoeDsaConfig(DeepseekV32Config): "layers.*.self_attn.q_b_proj": "colwise", "layers.*.self_attn.kv_b_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py index cfb81a96227d..09e98780ca9d 100755 --- a/src/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -395,6 +395,7 @@ def set_input_embeddings(self, value): class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} _tp_plan = {"embed_out": "colwise_allgather"} + _tp_ep_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_neox/modular_gpt_neox.py b/src/transformers/models/gpt_neox/modular_gpt_neox.py index cabda14021f0..2d8d31deb650 100644 --- a/src/transformers/models/gpt_neox/modular_gpt_neox.py +++ b/src/transformers/models/gpt_neox/modular_gpt_neox.py @@ -342,6 +342,7 @@ def forward( class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"} _tp_plan = {"embed_out": "colwise_allgather"} + _tp_ep_plan = {"embed_out": "colwise_allgather"} _pp_plan = {"embed_out": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index ebb980bd8a9b..a26aaf567b9f 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -594,6 +594,8 @@ class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/granite/modeling_granite.py b/src/transformers/models/granite/modeling_granite.py index 5a266b3acc93..e9fad65eabd5 100644 --- a/src/transformers/models/granite/modeling_granite.py +++ b/src/transformers/models/granite/modeling_granite.py @@ -447,6 +447,8 @@ class GraniteForCausalLM(GranitePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index 4f79272f164f..dae4348efbcf 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -628,6 +628,8 @@ class GraniteMoeForCausalLM(GraniteMoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 929a82f5ed82..fe927faa79d0 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1309,6 +1309,8 @@ class GraniteMoeHybridForCausalLM(GraniteMoeHybridPreTrainedModel, GenerationMix _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 8d0bc797a57b..b6f2884fec46 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -697,6 +697,8 @@ class GraniteMoeSharedForCausalLM(GraniteMoeSharedPreTrainedModel, GenerationMix _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/helium/modeling_helium.py b/src/transformers/models/helium/modeling_helium.py index b9a0e1121809..6c5a91b0725e 100644 --- a/src/transformers/models/helium/modeling_helium.py +++ b/src/transformers/models/helium/modeling_helium.py @@ -423,6 +423,8 @@ class HeliumForCausalLM(HeliumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/hrm_text/modeling_hrm_text.py b/src/transformers/models/hrm_text/modeling_hrm_text.py index 731f3cba1ebb..24d757afd551 100644 --- a/src/transformers/models/hrm_text/modeling_hrm_text.py +++ b/src/transformers/models/hrm_text/modeling_hrm_text.py @@ -557,6 +557,8 @@ class HrmTextForCausalLM(HrmTextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py index 3b2386cf8ab3..88e79b690295 100644 --- a/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py +++ b/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py @@ -463,6 +463,8 @@ class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py index db33809260af..7eea7af92237 100644 --- a/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py +++ b/src/transformers/models/hunyuan_v1_moe/modeling_hunyuan_v1_moe.py @@ -552,6 +552,8 @@ class HunYuanMoEV1ForCausalLM(HunYuanMoEV1PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py index dddb7269f95e..7335ba8d7a78 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -57,8 +57,8 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_norm": "activation_seq_dim_2", "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.self_attn.o_proj": "rowwise_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/hy_v3/modeling_hy_v3.py b/src/transformers/models/hy_v3/modeling_hy_v3.py index 6ade5c2d278d..c335daf46af1 100644 --- a/src/transformers/models/hy_v3/modeling_hy_v3.py +++ b/src/transformers/models/hy_v3/modeling_hy_v3.py @@ -546,6 +546,8 @@ class HYV3ForCausalLM(HYV3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py index 5c9547a2fe0e..c109e8973aa5 100644 --- a/src/transformers/models/hy_v3/modular_hy_v3.py +++ b/src/transformers/models/hy_v3/modular_hy_v3.py @@ -82,8 +82,8 @@ class HYV3Config(PreTrainedConfig): "layers.*.self_attn.q_norm": "activation_seq_dim_2", "layers.*.self_attn.k_norm": "activation_seq_dim_2", "layers.*.self_attn.o_proj": "rowwise_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/hyperclovax/modeling_hyperclovax.py b/src/transformers/models/hyperclovax/modeling_hyperclovax.py index a69fb46b121a..a5420a95e290 100644 --- a/src/transformers/models/hyperclovax/modeling_hyperclovax.py +++ b/src/transformers/models/hyperclovax/modeling_hyperclovax.py @@ -454,6 +454,8 @@ class HyperCLOVAXForCausalLM(HyperCLOVAXPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/jais2/modeling_jais2.py b/src/transformers/models/jais2/modeling_jais2.py index 1fb6c19d9044..b74f8acca655 100644 --- a/src/transformers/models/jais2/modeling_jais2.py +++ b/src/transformers/models/jais2/modeling_jais2.py @@ -399,6 +399,8 @@ class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index ba5a759c4636..ae303a3aec07 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -845,6 +845,8 @@ class JambaForCausalLM(JambaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py index fbe51b44259d..22dd63b3a71d 100644 --- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py +++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py @@ -876,6 +876,8 @@ class KyutaiSpeechToTextForConditionalGeneration(KyutaiSpeechToTextPreTrainedMod _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} _keep_in_fp32_modules_strict = ["codec_model"] diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py index 50e011ef08be..131444a506c4 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -67,8 +67,34 @@ class LagunaConfig(PreTrainedConfig): "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.g_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/laguna/modeling_laguna.py b/src/transformers/models/laguna/modeling_laguna.py index 61ffc6388ff7..d43737ef6282 100644 --- a/src/transformers/models/laguna/modeling_laguna.py +++ b/src/transformers/models/laguna/modeling_laguna.py @@ -673,6 +673,8 @@ class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/laguna/modular_laguna.py b/src/transformers/models/laguna/modular_laguna.py index ca65fe60dfad..00ddb037c721 100644 --- a/src/transformers/models/laguna/modular_laguna.py +++ b/src/transformers/models/laguna/modular_laguna.py @@ -85,8 +85,27 @@ class LagunaConfig(Qwen2MoeConfig): "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.g_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/lfm2/modeling_lfm2.py b/src/transformers/models/lfm2/modeling_lfm2.py index 73a88cef3734..5e75dbed8def 100644 --- a/src/transformers/models/lfm2/modeling_lfm2.py +++ b/src/transformers/models/lfm2/modeling_lfm2.py @@ -561,6 +561,8 @@ class Lfm2ForCausalLM(Lfm2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py index 0a0b19a11754..2e629902da46 100644 --- a/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -651,6 +651,8 @@ class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py index 655e542358f5..7cd78fce37f5 100644 --- a/src/transformers/models/llama/modeling_llama.py +++ b/src/transformers/models/llama/modeling_llama.py @@ -430,6 +430,8 @@ class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 781de8a57fd0..2612a236d0eb 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -128,11 +128,28 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.self_attn.o_proj": "rowwise_allreduce", "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear + "layers.*.feed_forward.experts": "moe_experts_allreduce", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", "layers.*.feed_forward.down_proj": "rowwise_allreduce", "layers.*.feed_forward.router": "ep_router", } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.feed_forward.shared_expert.gate_proj": "colwise", + "layers.*.feed_forward.shared_expert.up_proj": "colwise", + "layers.*.feed_forward.shared_expert.down_proj": "rowwise_allreduce", + "layers.*.feed_forward.gate_proj": "colwise", + "layers.*.feed_forward.up_proj": "colwise", + "layers.*.feed_forward.down_proj": "rowwise_allreduce", + "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", + "layers.*.feed_forward.experts.down_proj": "grouped_gemm", + "layers.*.feed_forward.experts": "moe_experts_allreduce", + "layers.*.feed_forward.router": "ep_router", + } vocab_size: int = 202048 hidden_size: int = 5120 diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index f3bdeda372fb..01ede0ab91c1 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -593,6 +593,8 @@ class Llama4ForCausalLM(Llama4PreTrainedModel, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} config: Llama4TextConfig def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/models/longcat_flash/configuration_longcat_flash.py b/src/transformers/models/longcat_flash/configuration_longcat_flash.py index ffeb6de9d01f..45fa53fc824d 100644 --- a/src/transformers/models/longcat_flash/configuration_longcat_flash.py +++ b/src/transformers/models/longcat_flash/configuration_longcat_flash.py @@ -57,6 +57,10 @@ class LongcatFlashConfig(PreTrainedConfig): "layers.*.self_attn.*.q_b_proj": "colwise", "layers.*.self_attn.*.kv_b_proj": "colwise", "layers.*.self_attn.*.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts.identity_expert": "moe_identity_expert", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlps.*.gate_proj": "colwise", "layers.*.mlps.*.up_proj": "colwise", "layers.*.mlps.*.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/longcat_flash/modeling_longcat_flash.py b/src/transformers/models/longcat_flash/modeling_longcat_flash.py index ae1084b1c1a7..b4fe7f92c4f7 100644 --- a/src/transformers/models/longcat_flash/modeling_longcat_flash.py +++ b/src/transformers/models/longcat_flash/modeling_longcat_flash.py @@ -646,6 +646,8 @@ class LongcatFlashForCausalLM(LongcatFlashPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} _keys_to_ignore_on_load_unexpected = [r"model\.mtp.*"] diff --git a/src/transformers/models/mellum/configuration_mellum.py b/src/transformers/models/mellum/configuration_mellum.py index f29681a631f3..4da68e0cf56c 100644 --- a/src/transformers/models/mellum/configuration_mellum.py +++ b/src/transformers/models/mellum/configuration_mellum.py @@ -55,6 +55,8 @@ class MellumConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", @@ -84,6 +86,42 @@ class MellumConfig(PreTrainedConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -137,13 +175,21 @@ def __post_init__(self, **kwargs): "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, } - self._update_sp_plan() + self._update_parallel_plans() super().__post_init__( **kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}, ) + def _is_moe_layer(self, layer_idx: int) -> bool: + return self.mlp_layer_types[layer_idx] == "sparse" + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + def _update_sp_plan(self): """Set per-layer SP entries depending on whether the MLP is dense or sparse.""" self.base_model_sp_plan = self.base_model_sp_plan.copy() @@ -161,6 +207,52 @@ def _update_sp_plan(self): self.base_model_sp_plan.update( { f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", f"layers.{i}.mlp.experts": "moe_experts_allreduce", } ) diff --git a/src/transformers/models/mellum/modeling_mellum.py b/src/transformers/models/mellum/modeling_mellum.py index 1b4ba4f1d725..db04cfa0ac9a 100644 --- a/src/transformers/models/mellum/modeling_mellum.py +++ b/src/transformers/models/mellum/modeling_mellum.py @@ -635,6 +635,8 @@ class MellumForCausalLM(MellumPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/mellum/modular_mellum.py b/src/transformers/models/mellum/modular_mellum.py index 2fb71f4ebfa8..da08ca93e687 100644 --- a/src/transformers/models/mellum/modular_mellum.py +++ b/src/transformers/models/mellum/modular_mellum.py @@ -68,6 +68,42 @@ class MellumConfig(Qwen3MoeConfig): "layers.*.post_attention_layernorm": "activation", "norm": "activation", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } vocab_size: int = 98304 hidden_size: int = 2304 @@ -99,7 +135,7 @@ def __post_init__(self, **kwargs): "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, } - self._update_sp_plan() + self._update_parallel_plans() PreTrainedConfig.__post_init__( self, @@ -107,6 +143,14 @@ def __post_init__(self, **kwargs): ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}, ) + def _is_moe_layer(self, layer_idx: int) -> bool: + return self.mlp_layer_types[layer_idx] == "sparse" + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + def _update_sp_plan(self): """Set per-layer SP entries depending on whether the MLP is dense or sparse.""" self.base_model_sp_plan = self.base_model_sp_plan.copy() @@ -124,6 +168,52 @@ def _update_sp_plan(self): self.base_model_sp_plan.update( { f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + for i, mlp_type in enumerate(self.mlp_layer_types): + if mlp_type == "dense": + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", f"layers.{i}.mlp.experts": "moe_experts_allreduce", } ) diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index d60e6d0c2849..8ff7d9a6d3db 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -66,6 +66,8 @@ class MiniMaxConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 7da5e5867ce2..0f6392c9976a 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -791,6 +791,8 @@ class MiniMaxForCausalLM(MiniMaxPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 8442eebd53e7..6f9ec2655a58 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -92,6 +92,8 @@ class MiniMaxConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 6193523d95b3..9adab09b43cf 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -52,6 +52,8 @@ class MiniMaxM2Config(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise_allgather", "layers.*.self_attn.v_proj": "colwise_allgather", "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_sp_plan = { @@ -64,6 +66,8 @@ class MiniMaxM2Config(PreTrainedConfig): "layers.*.self_attn.o_proj": "vocab_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index bed74f93540f..b14efdb3ceb5 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -589,6 +589,8 @@ class MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index 924fffe41477..b5917b630669 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -71,6 +71,8 @@ class MiniMaxM2Config(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise_allgather", "layers.*.self_attn.v_proj": "colwise_allgather", "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_sp_plan = { @@ -83,6 +85,8 @@ class MiniMaxM2Config(PreTrainedConfig): "layers.*.self_attn.o_proj": "vocab_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } diff --git a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py index fed287364099..41ba63ba5eaa 100644 --- a/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/modular_minimax_m3_vl.py @@ -731,7 +731,9 @@ class MiniMaxM3VLForCausalLM(MiniMaxM2ForCausalLM): config: MiniMaxM3VLTextConfig _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} + _tp_ep_plan = AttributeError() _sp_plan = AttributeError() + _sp_ep_plan = AttributeError() _fsdp_plan = AttributeError() def __init__(self, config: MiniMaxM3VLTextConfig): diff --git a/src/transformers/models/ministral/modeling_ministral.py b/src/transformers/models/ministral/modeling_ministral.py index d2faea85b8cc..0e1d344bf25f 100644 --- a/src/transformers/models/ministral/modeling_ministral.py +++ b/src/transformers/models/ministral/modeling_ministral.py @@ -431,6 +431,8 @@ class MinistralForCausalLM(MinistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/ministral3/modeling_ministral3.py b/src/transformers/models/ministral3/modeling_ministral3.py index 066a71997913..8b7dfec23fbe 100644 --- a/src/transformers/models/ministral3/modeling_ministral3.py +++ b/src/transformers/models/ministral3/modeling_ministral3.py @@ -414,6 +414,8 @@ class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py index d94ff3d6a312..5fa28011ccfa 100644 --- a/src/transformers/models/mistral/modeling_mistral.py +++ b/src/transformers/models/mistral/modeling_mistral.py @@ -403,6 +403,8 @@ class MistralForCausalLM(MistralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index c8218ffa5c57..9979792faffe 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -47,6 +47,8 @@ class Mistral4Config(PreTrainedConfig): model_type = "mistral4" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_experts.gate_proj": "colwise", "layers.*.mlp.shared_experts.up_proj": "colwise", diff --git a/src/transformers/models/mistral4/modeling_mistral4.py b/src/transformers/models/mistral4/modeling_mistral4.py index 6b62c4cbee09..10eecfa312ea 100644 --- a/src/transformers/models/mistral4/modeling_mistral4.py +++ b/src/transformers/models/mistral4/modeling_mistral4.py @@ -643,6 +643,8 @@ class Mistral4ForCausalLM(Mistral4PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index 41beb623719c..717e672afe00 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -48,6 +48,8 @@ class MixtralConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } @@ -62,6 +64,43 @@ class MixtralConfig(PreTrainedConfig): "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Inference TP + EP (complete plan — no runtime merge). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (complete plan — no runtime merge). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index df71b44ebac2..b0947d3f4271 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -582,6 +582,8 @@ class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/nanochat/modeling_nanochat.py b/src/transformers/models/nanochat/modeling_nanochat.py index a191eb2d5425..48011dac3736 100644 --- a/src/transformers/models/nanochat/modeling_nanochat.py +++ b/src/transformers/models/nanochat/modeling_nanochat.py @@ -434,6 +434,8 @@ class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/nanochat/modular_nanochat.py b/src/transformers/models/nanochat/modular_nanochat.py index 460cf321150f..6c6f79ae5750 100644 --- a/src/transformers/models/nanochat/modular_nanochat.py +++ b/src/transformers/models/nanochat/modular_nanochat.py @@ -201,6 +201,8 @@ class NanoChatForCausalLM(Gemma2ForCausalLM): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs) -> CausalLMOutputWithPast: r""" diff --git a/src/transformers/models/olmo/modeling_olmo.py b/src/transformers/models/olmo/modeling_olmo.py index 41169f2f389c..df5a149eaf91 100644 --- a/src/transformers/models/olmo/modeling_olmo.py +++ b/src/transformers/models/olmo/modeling_olmo.py @@ -428,6 +428,8 @@ class OlmoForCausalLM(OlmoPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/olmo2/modeling_olmo2.py b/src/transformers/models/olmo2/modeling_olmo2.py index c41d78cbcb9d..b03d73d36929 100644 --- a/src/transformers/models/olmo2/modeling_olmo2.py +++ b/src/transformers/models/olmo2/modeling_olmo2.py @@ -432,6 +432,8 @@ class Olmo2ForCausalLM(Olmo2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/olmo3/modeling_olmo3.py b/src/transformers/models/olmo3/modeling_olmo3.py index f5e41143fba6..4d765e7d56bd 100644 --- a/src/transformers/models/olmo3/modeling_olmo3.py +++ b/src/transformers/models/olmo3/modeling_olmo3.py @@ -436,6 +436,8 @@ class Olmo3ForCausalLM(Olmo3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py index c8a202a82f99..250c2f65b425 100644 --- a/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py +++ b/src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py @@ -1048,6 +1048,8 @@ class OlmoHybridForCausalLM(OlmoHybridPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index feedb02bfda4..01cc17a429d7 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -50,6 +50,8 @@ class OlmoeConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise_allgather", # due to the norm, we have to gather "layers.*.self_attn.v_proj": "colwise_allgather", # due to the norm, we have to gather "layers.*.self_attn.o_proj": "vocab_allreduce", # due to the norm, we have to gather + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_sp_plan = { @@ -62,6 +64,41 @@ class OlmoeConfig(PreTrainedConfig): "layers.*.self_attn.o_proj": "vocab_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "norm": "activation", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise_allgather", + "layers.*.self_attn.k_proj": "colwise_allgather", + "layers.*.self_attn.v_proj": "colwise_allgather", + "layers.*.self_attn.o_proj": "vocab_reduce_scatter", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index 10950656c462..d25bdd6f5282 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -606,6 +606,8 @@ class OlmoeForCausalLM(OlmoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py index 0acf6a8060aa..12317fbfc3cd 100644 --- a/src/transformers/models/phi/modeling_phi.py +++ b/src/transformers/models/phi/modeling_phi.py @@ -408,6 +408,8 @@ class PhiForCausalLM(PhiPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/phi3/modeling_phi3.py b/src/transformers/models/phi3/modeling_phi3.py index f0d6726b73d9..7ea8ccfcb241 100644 --- a/src/transformers/models/phi3/modeling_phi3.py +++ b/src/transformers/models/phi3/modeling_phi3.py @@ -434,6 +434,8 @@ class Phi3ForCausalLM(Phi3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 383445a29f8b..99003339fed4 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -1598,6 +1598,8 @@ class Phi4MultimodalForCausalLM(Phi4MultimodalPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index d933a6ec69d9..753b2535c1e4 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -774,6 +774,8 @@ class PhimoeForCausalLM(PhimoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/pi0/modeling_pi0.py b/src/transformers/models/pi0/modeling_pi0.py index 55ca4ca3f1be..ff3055f47d29 100644 --- a/src/transformers/models/pi0/modeling_pi0.py +++ b/src/transformers/models/pi0/modeling_pi0.py @@ -220,6 +220,7 @@ class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" _tp_plan = {"action_out_proj": "colwise_allgather"} + _tp_ep_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/pi0/modular_pi0.py b/src/transformers/models/pi0/modular_pi0.py index 51da5eff9629..968844bfb7a4 100644 --- a/src/transformers/models/pi0/modular_pi0.py +++ b/src/transformers/models/pi0/modular_pi0.py @@ -480,6 +480,7 @@ class PI0ForConditionalGeneration(PI0PreTrainedModel): """PI0 model with action projection heads and flow matching.""" _tp_plan = {"action_out_proj": "colwise_allgather"} + _tp_ep_plan = {"action_out_proj": "colwise_allgather"} def __init__(self, config: PI0Config): super().__init__(config) diff --git a/src/transformers/models/qwen2/modeling_qwen2.py b/src/transformers/models/qwen2/modeling_qwen2.py index 7dcc274a2cf2..ded1903f822f 100644 --- a/src/transformers/models/qwen2/modeling_qwen2.py +++ b/src/transformers/models/qwen2/modeling_qwen2.py @@ -418,6 +418,8 @@ class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index 329a6ef2a058..b000ad0942f3 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -58,6 +58,29 @@ class Qwen2MoeConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + } + # Expert-only EP plan: shards MoE experts along the expert axis; gate stays replicated. + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise_allreduce", diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index f7142b4650fb..3303d94c4848 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -619,6 +619,8 @@ class Qwen2MoeForCausalLM(Qwen2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py index 9c5b65f47896..84d55ae84b47 100644 --- a/src/transformers/models/qwen2_moe/modular_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modular_qwen2_moe.py @@ -233,6 +233,8 @@ class Qwen2MoeForCausalLM(MixtralForCausalLM, GenerationMixin): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 2ba669225879..fe3de31fcac7 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -443,6 +443,8 @@ class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 5df1e4c1631a..ba55590dbb0e 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1614,6 +1614,8 @@ class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5TextConfig diff --git a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py index 9b34c759b4bf..ee125e10a6e9 100644 --- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -61,6 +61,8 @@ class Qwen3_5MoeTextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index a1617dac1599..c7a62e7101b6 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -1801,6 +1801,8 @@ class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config: Qwen3_5MoeTextConfig @@ -1911,6 +1913,8 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMi _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py index edbac33dc6ec..306ac4568ff9 100644 --- a/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py @@ -90,6 +90,8 @@ class Qwen3_5MoeTextConfig(Qwen3NextConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", @@ -249,6 +251,8 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): _tp_plan = {"lm_head": "colwise_allgather"} _fsdp_plan = {"lm_head": "keep_full_weight"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} def forward(self, **super_kwargs): r""" diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py index 4395b7c622bd..21893966542c 100644 --- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py @@ -58,6 +58,8 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", @@ -90,6 +92,42 @@ class Qwen3MoeConfig(PreTrainedConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -135,6 +173,86 @@ def __post_init__(self, **kwargs): self.sliding_window = self.sliding_window if self.use_sliding_window else None self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers super().__post_init__(**kwargs) + self._update_parallel_plans() + + def _is_moe_layer(self, layer_idx: int) -> bool: + return (layer_idx not in self.mlp_only_layers) and ( + self.num_experts > 0 and (layer_idx + 1) % self.decoder_sparse_step == 0 + ) + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + self.base_model_sp_plan = self.base_model_sp_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) __all__ = ["Qwen3MoeConfig"] diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index 012e0cac7243..4b71fa5bbc5b 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -611,6 +611,8 @@ class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index b499291d70e2..f09c965d1093 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -66,6 +66,8 @@ class Qwen3NextConfig(PreTrainedConfig): "layers.*.mlp.shared_expert.gate_proj": "colwise", "layers.*.mlp.shared_expert.up_proj": "colwise", "layers.*.mlp.shared_expert.down_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index a592ad3b7b43..5ef52f8227ba 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1089,6 +1089,8 @@ class Qwen3NextForCausalLM(Qwen3NextPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py index b18e65f61b74..d52c179d58ff 100644 --- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py @@ -396,6 +396,8 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", @@ -425,6 +427,42 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -467,8 +505,89 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): def __post_init__(self, **kwargs): self.sliding_window = self.sliding_window + self.sliding_window = self.sliding_window if self.use_sliding_window else None self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers super().__post_init__(**kwargs) + self._update_parallel_plans() + + def _is_moe_layer(self, layer_idx: int) -> bool: + return (layer_idx not in self.mlp_only_layers) and ( + self.num_experts > 0 and (layer_idx + 1) % self.decoder_sparse_step == 0 + ) + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + self.base_model_sp_plan = self.base_model_sp_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) @auto_docstring(checkpoint="Qwen/Qwen3-30B-A3B-Base") diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py index f4ea11cbb274..a6099aa8df41 100644 --- a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py @@ -2634,6 +2634,8 @@ class Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration(Qwen3OmniMoeP _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerCodePredictorConfig @@ -3020,6 +3022,8 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3OmniMoeThinkerTextPreTrain _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} _tp_plan = {"codec_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"codec_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerConfig diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index 3b0211193954..5d1ab7fe5a62 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -401,6 +401,42 @@ class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } vocab_size: int = 3072 hidden_size: int = 1024 @@ -413,8 +449,8 @@ class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): use_sliding_window = AttributeError() def __post_init__(self, **kwargs): - super().__post_init__(**kwargs) self.sliding_window = self.sliding_window + super().__post_init__(**kwargs) @auto_docstring(checkpoint="Qwen/Qwen3-30B-A3B-Base") @@ -1665,7 +1701,11 @@ def get_input_embeddings(self): class Qwen3OmniMoeTalkerForConditionalGeneration(Qwen3MoeForCausalLM): _tied_weights_keys = {"codec_head": "model.codec_embedding.weight"} _tp_plan = {"codec_head": "colwise_allgather"} + _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"codec_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"codec_head": (["hidden_states"], ["logits"])} + _fsdp_plan = {"lm_head": "keep_full_weight"} config_class = Qwen3OmniMoeTalkerConfig base_model_prefix = "talker" _no_split_modules = ["Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration"] diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index f5e2749d473e..00eb23290288 100644 --- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py @@ -88,6 +88,42 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -136,8 +172,89 @@ def __post_init__(self, **kwargs): self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads self.sliding_window = None + self.sliding_window = self.sliding_window if self.use_sliding_window else None self.mlp_only_layers = [] if self.mlp_only_layers is None else self.mlp_only_layers super().__post_init__(**kwargs) + self._update_parallel_plans() + + def _is_moe_layer(self, layer_idx: int) -> bool: + return (layer_idx not in self.mlp_only_layers) and ( + self.num_experts > 0 and (layer_idx + 1) % self.decoder_sparse_step == 0 + ) + + def _update_parallel_plans(self): + self._update_sp_plan() + self._update_sp_ep_plan() + self._update_tp_ep_plan() + + def _update_sp_plan(self): + self.base_model_sp_plan = self.base_model_sp_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + f"layers.{i}.mlp.experts.down_proj": "moe_tp_down_rowwise", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_sp_ep_plan(self): + self.base_model_sp_ep_plan = self.base_model_sp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather_split", + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_sp_ep_plan.update( + { + f"layers.{i}.mlp": "module_allgather", + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_reduce_scatter", + } + ) + + def _update_tp_ep_plan(self): + self.base_model_tp_ep_plan = self.base_model_tp_ep_plan.copy() + + for i in range(self.num_hidden_layers): + if self._is_moe_layer(i): + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate": "ep_router", + f"layers.{i}.mlp.experts.gate_up_proj": "grouped_gemm", + f"layers.{i}.mlp.experts.down_proj": "grouped_gemm", + f"layers.{i}.mlp.experts": "moe_experts_allreduce", + } + ) + else: + self.base_model_tp_ep_plan.update( + { + f"layers.{i}.mlp.gate_proj": "colwise", + f"layers.{i}.mlp.up_proj": "colwise", + f"layers.{i}.mlp.down_proj": "rowwise_allreduce", + } + ) @auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py index 555e3d63ccb3..91dbf4b3241d 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -97,6 +97,42 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig): "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce", } + # Inference TP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_tp_ep_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_allreduce", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + } + # Training SP + EP (per-layer overrides applied in `_update_parallel_plans`). + base_model_sp_ep_plan = { + "embed_tokens": "vocab_reduce_scatter", + "layers.*.input_layernorm": "activation", + "layers.*.self_attn": "module_allgather_hidden_states", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", + "layers.*.self_attn.q_norm": "activation_seq_dim_2", + "layers.*.self_attn.k_norm": "activation_seq_dim_2", + "layers.*.post_attention_layernorm": "activation", + "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_experts_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise_reduce_scatter", + "norm": "activation", + } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), @@ -125,8 +161,8 @@ def __post_init__(self, **kwargs): self.num_key_value_heads = self.num_attention_heads self.head_dim = self.head_dim or self.hidden_size // self.num_attention_heads - super().__post_init__(**kwargs) self.sliding_window = None + super().__post_init__(**kwargs) @auto_docstring(checkpoint="Qwen/Qwen3-VL-30B-A3B-Instruct") diff --git a/src/transformers/models/seed_oss/modeling_seed_oss.py b/src/transformers/models/seed_oss/modeling_seed_oss.py index 841765a7c928..e2f667e03a36 100644 --- a/src/transformers/models/seed_oss/modeling_seed_oss.py +++ b/src/transformers/models/seed_oss/modeling_seed_oss.py @@ -431,6 +431,8 @@ class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/smollm3/modeling_smollm3.py b/src/transformers/models/smollm3/modeling_smollm3.py index d82882382bb9..a8077ced3274 100644 --- a/src/transformers/models/smollm3/modeling_smollm3.py +++ b/src/transformers/models/smollm3/modeling_smollm3.py @@ -447,6 +447,8 @@ class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index 5d9aec3ac7a0..aefa196808fe 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -42,6 +42,8 @@ class SolarOpenConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_pp_plan = { @@ -96,6 +98,8 @@ class SolarOpenConfig(PreTrainedConfig): "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } diff --git a/src/transformers/models/solar_open/modeling_solar_open.py b/src/transformers/models/solar_open/modeling_solar_open.py index 53abac7c7d07..3ba4782d3439 100644 --- a/src/transformers/models/solar_open/modeling_solar_open.py +++ b/src/transformers/models/solar_open/modeling_solar_open.py @@ -554,6 +554,8 @@ class SolarOpenForCausalLM(SolarOpenPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/solar_open/modular_solar_open.py b/src/transformers/models/solar_open/modular_solar_open.py index 48bba9439e08..51cf4c0e0b39 100644 --- a/src/transformers/models/solar_open/modular_solar_open.py +++ b/src/transformers/models/solar_open/modular_solar_open.py @@ -48,6 +48,8 @@ class SolarOpenConfig(Glm4MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise_allreduce", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", } base_model_sp_plan = { @@ -60,6 +62,8 @@ class SolarOpenConfig(Glm4MoeConfig): "layers.*.self_attn.o_proj": "rowwise_reduce_scatter", "layers.*.post_attention_layernorm": "activation", "layers.*.mlp": "module_allgather_split", + "layers.*.mlp.experts.gate_up_proj": "moe_tp_gate_up_colwise", + "layers.*.mlp.experts.down_proj": "moe_tp_down_rowwise", "layers.*.mlp.experts": "moe_experts_allreduce", "norm": "activation", } diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index baabf385da59..ab9ed58cfdfd 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -411,6 +411,8 @@ class Starcoder2ForCausalLM(Starcoder2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index ceafdfa43d98..c8aed75038a2 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -947,6 +947,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} _tp_plan = {"lm_head.out_proj": "colwise_allgather"} + _tp_ep_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py index f6a6595f9dbc..43fed4ea81a0 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -792,6 +792,7 @@ def forward( class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.out_proj.weight": "model.decoder.embed_tokens.weight"} _tp_plan = {"lm_head.out_proj": "colwise_allgather"} + _tp_ep_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5GemmaConfig): diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 830785694f1a..587cd8768420 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -1184,6 +1184,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } _tp_plan = {"lm_head.out_proj": "colwise_allgather"} + _tp_ep_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index 58384929d29e..10f741252cd0 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -981,6 +981,7 @@ class T5Gemma2ForConditionalGeneration(T5Gemma2PreTrainedModel, GenerationMixin) "lm_head.out_proj.weight": "model.encoder.text_model.embed_tokens.weight", } _tp_plan = {"lm_head.out_proj": "colwise_allgather"} + _tp_ep_plan = {"lm_head.out_proj": "colwise_allgather"} _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])} def __init__(self, config: T5Gemma2Config): diff --git a/src/transformers/models/vaultgemma/modeling_vaultgemma.py b/src/transformers/models/vaultgemma/modeling_vaultgemma.py index 1ddba4af2210..72a1b88beaae 100644 --- a/src/transformers/models/vaultgemma/modeling_vaultgemma.py +++ b/src/transformers/models/vaultgemma/modeling_vaultgemma.py @@ -469,6 +469,8 @@ class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/src/transformers/models/youtu/modeling_youtu.py b/src/transformers/models/youtu/modeling_youtu.py index 5d64e08f0256..f2e2211337a8 100644 --- a/src/transformers/models/youtu/modeling_youtu.py +++ b/src/transformers/models/youtu/modeling_youtu.py @@ -536,6 +536,8 @@ class YoutuForCausalLM(YoutuPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_allgather"} _sp_plan = {"lm_head": "colwise_loss_parallel"} + _tp_ep_plan = {"lm_head": "colwise_allgather"} + _sp_ep_plan = {"lm_head": "colwise_loss_parallel"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _fsdp_plan = {"lm_head": "keep_full_weight"} diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index fa39405791f4..a2c0f4097c07 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -68,6 +68,7 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase): _is_stateful = True model_split_percents = [0.5, 0.6] model_tester_class = GptOssModelTester + tensor_parallel_atol = 7e-3 @require_kernels def test_kernelize_does_not_crash(self): diff --git a/tests/models/hrm_text/test_modeling_hrm_text.py b/tests/models/hrm_text/test_modeling_hrm_text.py index 97644dede00d..441a632699b1 100644 --- a/tests/models/hrm_text/test_modeling_hrm_text.py +++ b/tests/models/hrm_text/test_modeling_hrm_text.py @@ -73,6 +73,14 @@ def test_tp_forward(self): def test_tp_backward(self): pass + @unittest.skip(reason="Higher tols (likely due to different recursion and grad patterns). FIXME later") + def test_sp_forward(self): + pass + + @unittest.skip(reason="Higher tols (likely due to different recursion and grad patterns). FIXME later") + def test_sp_backward(self): + pass + @unittest.skip(reason="Higher tols (likely due to different recursion and grad patterns). FIXME later") def test_tp_generation(self): pass diff --git a/tests/models/qwen3_moe/test_modeling_qwen3_moe.py b/tests/models/qwen3_moe/test_modeling_qwen3_moe.py index 13fafa6969ae..78cbe6b08ef7 100644 --- a/tests/models/qwen3_moe/test_modeling_qwen3_moe.py +++ b/tests/models/qwen3_moe/test_modeling_qwen3_moe.py @@ -45,6 +45,15 @@ class Qwen3MoeModelTester(CausalLMModelTester): if is_torch_available(): base_model_class = Qwen3MoeModel + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Make the tiny config build alternating dense/sparse layers so the SP plan is exercised + # against both a dense Qwen3MoeMLP and a sparse Qwen3MoeSparseMoeBlock. decoder_sparse_step=2 + # makes even layers dense and odd layers MoE; pin >=2 layers so a sparse layer actually + # exists (don't rely on the shared tester num_hidden_layers default staying >= 2). + self.num_hidden_layers = max(self.num_hidden_layers, 2) + self.decoder_sparse_step = 2 + @require_torch class Qwen3MoeModelTest(CausalLMModelTest, unittest.TestCase): diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index d71c398c039e..f35fa305ef44 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -62,31 +62,6 @@ def _find_free_port(): return s.getsockname()[1] -def get_packed_grad_shard(grad, world_size, rank, dim): - """Get the correct shard of a packed gradient (matching get_packed_weights interleaved logic). - - Packed weights like gate_up_proj are sharded with interleaving: - Original: [G0 G1 G2 G3 | U0 U1 U2 U3] (gate | up) - Rank 0: [G0 G1 | U0 U1] - Rank 1: [G2 G3 | U2 U3] - """ - total_size = grad.shape[dim] - # Packed weights have 2 blocks (gate and up) - block_size = total_size // 2 - shard_block_size = block_size // world_size - - # Build interleaved indices - indices = [] - for block_idx in range(2): # gate block, then up block - block_offset = block_idx * block_size - start = block_offset + rank * shard_block_size - stop = block_offset + (rank + 1) * shard_block_size - indices.extend(range(start, stop)) - - # Select along the sharded dimension - return grad.index_select(dim, torch.tensor(indices, device=grad.device)) - - def _global_wrapper(rank, func, tp, port, backend, func_args, func_kwargs): """Wrapper to set up distributed environment and run the test function.""" @@ -130,176 +105,164 @@ def wrapper(*args, **kwargs): return _init_distributed_inner -def _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=False): - """Load TP model and non-TP reference model for comparison. - - Returns: - tuple: (model_tp, model_ref, device) - """ +def _load_distributed_and_reference_models(model_path, model_class, mode: str): + """Load a distributed model and an unsharded reference model for comparison.""" tp_size = dist.get_world_size() - distributed_config = DistributedConfig(tp_size=tp_size, enable_sequence_parallel=enable_sequence_parallel) - model_tp = model_class.from_pretrained( - model_path, distributed_config=distributed_config, attn_implementation="sdpa" - ) + if mode in ("ep", "tp_ep"): + distributed_config = DistributedConfig(tp_size=tp_size, enable_expert_parallel=True) + elif mode in ("sp", "sp_ep"): + distributed_config = DistributedConfig( + tp_size=tp_size, enable_sequence_parallel=True, enable_expert_parallel=(mode == "sp_ep") + ) + else: + distributed_config = DistributedConfig(tp_size=tp_size, enable_sequence_parallel=False) + + from_pretrained_kwargs = {"distributed_config": distributed_config} + if mode not in ("ep", "tp_ep"): + from_pretrained_kwargs["attn_implementation"] = "sdpa" + + model_dist = model_class.from_pretrained(model_path, **from_pretrained_kwargs) dist.barrier() - device = model_tp.device - model_ref = model_class.from_pretrained(model_path, attn_implementation="sdpa") + device = model_dist.device + if mode in ("ep", "tp_ep"): + model_ref = model_class.from_pretrained(model_path) + else: + model_ref = model_class.from_pretrained(model_path, attn_implementation="sdpa") model_ref = model_ref.to(device) - return model_tp, model_ref, device - - -def _get_active_tp_plan(model_tp): - distributed_config = getattr(model_tp.config, "distributed_config", None) - tp_plan = getattr(distributed_config, "tp_plan", None) - return tp_plan or getattr(model_tp, "_tp_plan", None) or {} + return model_dist, model_ref, device -def _verify_tp_sharding(rank, model_tp, model_ref): - """Verify TP sharding by comparing parameter shapes between TP and reference models. - - For DTensor params, uses the local tensor shape (not the global DTensor shape). - - Returns: - list: Names of sharded parameters - """ - world_size = dist.get_world_size() - sharded_params = [] - tp_plan = _get_active_tp_plan(model_tp) - - for (name, param), (_, param_full) in zip(model_tp.named_parameters(), model_ref.named_parameters()): - # For DTensor params, get the local shape for comparison - param_local = _to_local(param) - if param_local.shape != param_full.shape: - sharded_params.append(name) - if rank == 0: - print(f"[TP Test Debug] TP sharded: {name} - full: {param_full.shape} -> sharded: {param_local.shape}") - - # Verify sharding is correct - for dim in range(param_local.ndim): - if param_local.size(dim) != param_full.size(dim): - param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) - if param_plan == "packed_colwise": - expected_size = param_full.size(dim) // world_size - assert param_local.size(dim) == expected_size, ( - f"Packed weight {name} sharding incorrect: expected {expected_size}, got {param_local.size(dim)}" - ) - else: - expected_size = (param_full.size(dim) + world_size - 1) // world_size - assert param_local.size(dim) <= expected_size, ( - f"Weight {name} sharding incorrect: expected <= {expected_size}, got {param_local.size(dim)}" - ) - break - - return sharded_params - - -def _test_tp_forward_impl(_rank, model_path, model_class, atol, rtol): - """Implementation for comparing TP and non-TP model outputs.""" +def _test_forward_impl(_rank, model_path, model_class, atol, rtol, mode: str): + """Compare distributed and reference model forward outputs for the given mode.""" + assert mode in ("tp", "sp", "ep", "tp_ep", "sp_ep") set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) - - _verify_tp_sharding(_rank, model_tp, model) + model_dist, model_ref, device = _load_distributed_and_reference_models(model_path, model_class, mode) - model_tp.eval() - model.eval() + model_dist.eval() + model_ref.eval() - vocab_size = model.config.vocab_size + vocab_size = model_ref.config.vocab_size set_seed(0) input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) with torch.no_grad(): - logits = model(input_ids).logits - logits_tp = _to_local(model_tp(input_ids).logits) + logits_ref = model_ref(input_ids).logits + logits_dist = model_dist(input_ids).logits + if mode in ("tp", "sp", "sp_ep"): + logits_dist = _to_local(logits_dist) - diff = (logits - logits_tp).abs() - assert torch.allclose(logits, logits_tp, atol=atol, rtol=rtol), ( - f"TP and non-TP model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}" + diff = (logits_ref - logits_dist).abs() + assert torch.allclose(logits_ref, logits_dist, atol=atol, rtol=rtol), ( + f"{mode.upper()} and reference model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}" ) dist.barrier() -def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): - """Implementation for comparing TP and non-TP model backward passes.""" +def _get_packed_grad_shard(grad, world_size, rank, dim): + """Get the correct shard of a packed gradient (matching get_packed_weights interleaved logic). + + Packed weights (packed_colwise, moe_tp_gate_up_colwise) are sharded with interleaving: + Original: [G0 G1 G2 G3 | U0 U1 U2 U3] (gate | up) + Rank 0: [G0 G1 | U0 U1] + Rank 1: [G2 G3 | U2 U3] + """ + total_size = grad.shape[dim] + # Packed weights have 2 blocks (gate and up) + block_size = total_size // 2 + shard_block_size = block_size // world_size + + # Build interleaved indices + indices = [] + for block_idx in range(2): # gate block, then up block + block_offset = block_idx * block_size + start = block_offset + rank * shard_block_size + stop = block_offset + (rank + 1) * shard_block_size + indices.extend(range(start, stop)) + + # Select along the sharded dimension + return grad.index_select(dim, torch.tensor(indices, device=grad.device)) + + +def _test_backward_impl(rank, model_path, model_class, atol, rtol, mode: str): + """Compare distributed and reference model backward passes for the given mode.""" + assert mode in ("tp", "sp", "ep", "tp_ep", "sp_ep") set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class, enable_sequence_parallel=True) - tp_plan = _get_active_tp_plan(model_tp) - model_tp.train() - model.train() + model_dist, model_ref, device = _load_distributed_and_reference_models(model_path, model_class, mode) + applied_plan = getattr(model_dist, "tp_plan", None) or getattr(model_dist, f"_{mode}_plan", None) or {} + model_dist.train() + model_ref.train() - vocab_size = model.config.vocab_size + vocab_size = model_ref.config.vocab_size set_seed(0) input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) set_seed(0) labels = torch.randint(0, vocab_size, (2, 64)).to(device) - loss = model(input_ids, labels=labels, use_cache=False).loss - loss.backward() + loss_ref = model_ref(input_ids, labels=labels, use_cache=False).loss + loss_ref.backward() + + loss_dist = model_dist(input_ids, labels=labels, use_cache=False).loss + loss_dist.backward() - loss_tp = model_tp(input_ids, labels=labels, use_cache=False).loss - loss_tp.backward() + if mode in ("tp", "sp", "sp_ep"): + loss_dist_value = _to_local(loss_dist) + else: + loss_dist_value = loss_dist - loss_tp_local = _to_local(loss_tp) - assert torch.allclose(loss, loss_tp_local, atol=atol, rtol=rtol), ( - f"TP and non-TP model losses differ. " - f"Non-TP loss: {loss.item()}, TP loss: {loss_tp_local.item()}, " - f"Diff: {(loss - loss_tp_local).abs().item()}" + assert torch.allclose(loss_ref, loss_dist_value, atol=atol, rtol=rtol), ( + f"{mode.upper()} and reference model losses differ. " + f"Reference loss: {loss_ref.item()}, {mode.upper()} loss: {loss_dist_value.item()}, " + f"Diff: {(loss_ref - loss_dist_value).abs().item()}" ) - # Compare gradients for matching parameters + # Compare parameter gradients world_size = dist.get_world_size() + failed_grads = [] + for (name, param), (_, param_dist) in zip(model_ref.named_parameters(), model_dist.named_parameters()): + if param.grad is None or param_dist.grad is None: + continue + + grad = param.grad + grad_dist = _to_local(param_dist.grad) + + if grad.shape != grad_dist.shape: + for dim in range(grad.ndim): + if grad.size(dim) != grad_dist.size(dim): + param_plan = _get_parameter_tp_plan(name, applied_plan, is_weight=True) + if param_plan in ("packed_colwise", "moe_tp_gate_up_colwise"): + grad = _get_packed_grad_shard(grad, world_size, rank, dim) + else: + shard_size = grad_dist.size(dim) + start = rank * shard_size + grad = grad.narrow(dim, start, shard_size) + break - # Debug: check tied weights and parameter alignment - failed_grads = {} - for (name, param), (name_tp, param_tp) in zip(model.named_parameters(), model_tp.named_parameters()): - if param.grad is not None and param_tp.grad is not None: - grad = param.grad - grad_tp = _to_local(param_tp.grad) - - # Slice reference gradient to match local shard if parameter is sharded - if grad.shape != grad_tp.shape: - for dim in range(grad.ndim): - if grad.size(dim) != grad_tp.size(dim): - param_plan = _get_parameter_tp_plan(name, tp_plan, is_weight=True) - if param_plan == "packed_colwise": - # interleaved slicing - grad = get_packed_grad_shard(grad, world_size, rank, dim) - else: - # regular slicing - shard_size = grad_tp.size(dim) - start = rank * shard_size - grad = grad.narrow(dim, start, shard_size) - break - - if not torch.allclose(grad.cpu(), grad_tp.cpu(), atol=atol, rtol=rtol): - max_diff = (grad.cpu() - grad_tp.cpu()).abs().max().item() - ref_abs_max = grad.cpu().abs().max().item() - tp_abs_max = grad_tp.cpu().abs().max().item() - ratio = tp_abs_max / ref_abs_max if ref_abs_max > 0 else float("inf") - failed_grads[name] = (max_diff, ref_abs_max, tp_abs_max, ratio) - - assert not failed_grads, f"Gradients differ for {len(failed_grads)} parameter(s):\n" + "\n".join( - f" {name}: max_diff={v[0]:.6f}, ref_max={v[1]:.6f}, tp_max={v[2]:.6f}, tp/ref ratio={v[3]:.4f}" - for name, v in failed_grads.items() - ) + try: + torch.testing.assert_close(grad, grad_dist, atol=atol, rtol=rtol) + except AssertionError as e: + failed_grads.append(f"{name}: {e}") + + assert not failed_grads, "Gradients differ:\n" + "\n".join(failed_grads) dist.barrier() -def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new_tokens): - """Implementation for comparing TP and non-TP model generation outputs (direct load path).""" +def _test_generation_impl(_rank, model_path, model_class, atol, rtol, mode: str, max_new_tokens): + """Compare distributed and reference model generation outputs for the given mode.""" + assert mode in ("tp", "ep", "tp_ep") set_seed(0) - model_tp, model, device = _load_tp_and_reference_models(model_path, model_class) - model_tp.eval() - model.eval() + model_dist, model_ref, device = _load_distributed_and_reference_models(model_path, model_class, mode) + model_dist.eval() + model_ref.eval() set_seed(0) - vocab_size = model.config.vocab_size + vocab_size = model_ref.config.vocab_size input_ids = torch.randint(0, vocab_size, (1, 10)).to(device) generation_kwargs = { "max_new_tokens": max_new_tokens, @@ -311,38 +274,46 @@ def _test_tp_generation_impl(_rank, model_path, model_class, atol, rtol, max_new } with torch.no_grad(): - output = model.generate(input_ids, **generation_kwargs) - output_tp = model_tp.generate(input_ids, **generation_kwargs) + output = model_ref.generate(input_ids, **generation_kwargs) + output_dist = model_dist.generate(input_ids, **generation_kwargs) # Compare logits/scores at each generation step scores = torch.stack(output.scores) - scores_tp = torch.stack([_to_local(s) for s in output_tp.scores]) - - diff = (scores - scores_tp).abs() - assert torch.allclose(scores, scores_tp, atol=atol, rtol=rtol), ( - f"TP and non-TP model generation logits differ (direct load path). " + if mode in ("tp", "sp"): + scores_dist = torch.stack([_to_local(s) for s in output_dist.scores]) + else: + scores_dist = torch.stack(output_dist.scores) + + diff = (scores - scores_dist).abs() + assert torch.allclose(scores, scores_dist, atol=atol, rtol=rtol), ( + f"{mode.upper()} and reference model generation logits differ. " f"Max diff: {diff.max().item()} | Mean diff: {diff.mean().item()}" ) # Compare generated token sequences - sequences_tp = _to_local(output_tp.sequences) - assert torch.equal(output.sequences, sequences_tp), ( - f"TP and non-TP model generated different token sequences (direct load path). " - f"Non-TP: {output.sequences.tolist()} | TP: {sequences_tp.tolist()}" + sequences_dist = _to_local(output_dist.sequences) if mode in ("tp", "sp") else output_dist.sequences + assert torch.equal(output.sequences, sequences_dist), ( + f"{mode.upper()} and reference model generated different token sequences. " + f"Reference: {output.sequences.tolist()} | {mode.upper()}: {sequences_dist.tolist()}" ) dist.barrier() -def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_tokens): +def _test_generation_quantized_impl(_rank, model_path, model_class, max_new_tokens, mode: str = "tp"): """Implementation for comparing TP+quantized and non-TP quantized generation (sequence equality).""" + assert mode in ("tp", "tp_ep") set_seed(0) quantization_config = TorchAoConfig(Float8WeightOnlyConfig()) + distributed_config = DistributedConfig( + tp_size=dist.get_world_size(), + enable_expert_parallel=(mode == "tp_ep"), + ) model_tp = model_class.from_pretrained( model_path, - distributed_config=DistributedConfig(tp_size=dist.get_world_size()), + distributed_config=distributed_config, quantization_config=quantization_config, ) dist.barrier() @@ -392,76 +363,6 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t dist.barrier() -def _load_ep_and_reference_models(model_path, model_class): - """Load EP model and non-EP reference model for comparison.""" - tp_size = dist.get_world_size() - model_ep = model_class.from_pretrained( - model_path, - distributed_config=DistributedConfig( - tp_size=tp_size, - enable_expert_parallel=True, - ), - ) - dist.barrier() - - device = model_ep.device - model_ref = model_class.from_pretrained(model_path) - model_ref = model_ref.to(device) - - return model_ep, model_ref, device - - -def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol): - """Implementation for comparing EP and non-EP model outputs.""" - set_seed(0) - - model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class) - - model_ep.eval() - model_ref.eval() - - vocab_size = model_ref.config.vocab_size - input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) - - with torch.no_grad(): - logits_ref = model_ref(input_ids).logits - logits_ep = model_ep(input_ids).logits - - diff = (logits_ref - logits_ep).abs() - assert torch.allclose(logits_ref, logits_ep, atol=atol, rtol=rtol), ( - f"EP and non-EP model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}" - ) - - dist.barrier() - - -def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol): - """Implementation for comparing EP and non-EP model backward passes.""" - set_seed(0) - - model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class) - model_ep.train() - model_ref.train() - - vocab_size = model_ref.config.vocab_size - input_ids = torch.randint(0, vocab_size, (2, 64)).to(device) - labels = torch.randint(0, vocab_size, (2, 64)).to(device) - - loss_ref = model_ref(input_ids, labels=labels).loss - loss_ref.backward() - - loss_ep = model_ep(input_ids, labels=labels).loss - loss_ep.backward() - - assert torch.allclose(loss_ref, loss_ep, atol=atol, rtol=rtol), ( - f"EP and non-EP model losses differ. " - f"Non-EP loss: {loss_ref.item()}, EP loss: {loss_ep.item()}, " - f"Diff: {(loss_ref - loss_ep).abs().item()}" - ) - - dist.barrier() - - class TensorParallelTesterMixin(ABC): """ Mixin for tensor parallel tests. Add to model test classes alongside ModelTesterMixin. @@ -471,6 +372,13 @@ class TensorParallelTesterMixin(ABC): - causal_lm_class, base_model_class, etc. This mixin adds tensor parallel-specific tests using that infrastructure. + + Distributed test modes: + - "tp": _tp_plan, enable_sequence_parallel=False — inference-style TP. + - "sp": _sp_plan with per-layer MLP entries — training-style sequence parallel. + - "ep": _ep_plan, enable_expert_parallel=True — expert parallel on MoE weights (legacy EP-only path). + - "tp_ep": _tp_plan ∪ _ep_plan — inference TP dense + EP MoE. + - "sp_ep": _sp_plan ∪ _ep_plan — training SP dense + EP MoE. """ # ============================================================ @@ -499,14 +407,49 @@ def _has_tp_plan(self) -> bool: config = self.model_tester.get_config() return hasattr(config, "base_model_tp_plan") and config.base_model_tp_plan is not None - def _get_tp_model_class(self): + def _has_sp_plan(self) -> bool: + config = self.model_tester.get_config() + return getattr(config, "base_model_sp_plan", None) and config.base_model_sp_plan is not None + + def _get_model_class(self): """Get the model class to use for TP tests (prefers *ForCausalLM).""" if hasattr(self.model_tester, "causal_lm_class") and self.model_tester.causal_lm_class is not None: return self.model_tester.causal_lm_class return self.all_model_classes[0] - def _skip_if_not_supported(self): - """Check and skip test if TP is not supported for this model/environment.""" + def _assert_mixed_mlp_layers(self, model, config): + """If the config interleaves dense and sparse/MoE layers, fail fast unless the tiny test + model actually built both kinds. + """ + layer_types = getattr(config, "mlp_layer_types", None) + if layer_types is not None: + # MoE marker is 'sparse' or 'moe' depending on the model; dense is always 'dense'. + intends_mixed = "dense" in layer_types and any(t != "dense" for t in layer_types) + else: + sparse_step = getattr(config, "decoder_sparse_step", 1) or 1 + intends_mixed = sparse_step > 1 or bool(getattr(config, "mlp_only_layers", None)) + if not intends_mixed: + return + + kinds = { + "moe" if hasattr(module, "experts") else "dense" + for name, module in model.named_modules() + if name.endswith(".mlp") + } + # Empty means this model's feed-forward isn't named `.mlp` (e.g. `block_sparse_moe`); the + # structural check doesn't apply, so skip rather than spuriously fail its TP test. + if not kinds: + return + self.assertEqual( + kinds, + {"dense", "moe"}, + f"{type(model).__name__}: config interleaves dense/MoE mlp layers but the tiny test config " + f"built only {kinds}. Adjust the model tester (num_hidden_layers / decoder_sparse_step / " + f"mlp_layer_types) so both kinds exist.", + ) + + def _skip_if_mode_not_supported(self, mode: str) -> None: + """Check and skip test if the mode is not supported for this model/environment.""" if not is_torch_greater_or_equal("2.9"): self.skipTest("Tensor parallel tests require torch >= 2.9") @@ -522,60 +465,101 @@ def _skip_if_not_supported(self): if not hasattr(self.model_tester, "causal_lm_class") or self.model_tester.causal_lm_class is None: self.skipTest("Model tester does not have causal_lm_class (not using CausalLMModelTester)") - if not self._has_tp_plan(): + if mode == "tp" and not self._has_tp_plan(): self.skipTest("Model does not have a tensor parallel plan (base_model_tp_plan)") + elif mode == "sp" and not self._has_sp_plan(): + self.skipTest("Model does not have a sequence parallel plan (base_model_sp_plan)") + elif mode == "ep" and not self._has_ep_plan(): + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + elif mode == "tp_ep": + if not self._has_tp_plan(): + self.skipTest("Model does not have a tensor parallel plan (base_model_tp_plan)") + if not self._has_ep_plan(): + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + elif mode == "sp_ep": + if not self._has_sp_plan(): + self.skipTest("Model does not have a sequence parallel plan (base_model_sp_plan)") + if not self._has_ep_plan(): + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") # # Skip encoder-decoder models (TP not supported) # if getattr(self, "is_encoder_decoder", False): # self.skipTest("TP tests not supported for encoder-decoder models") - # # Skip VLM models for now - # config = self.model_tester.get_config() - # if hasattr(config, "vision_config") and config.vision_config is not None: - # self.skipTest("VLM models are not yet supported in TP tests") + @is_tensor_parallel_test + def test_tp_forward(self): + """Tensor parallel forward (_tp_plan, no sequence parallel).""" + self._skip_if_mode_not_supported("tp") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - def _skip_if_ep_not_supported(self): - """Check and skip test if EP is not supported for this model/environment.""" - if not is_torch_greater_or_equal("2.9"): - self.skipTest("Expert parallel tests require torch >= 2.9") + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)(tmp_dir, model_class, atol, rtol, "tp") - if torch.cuda.is_available() or torch.xpu.is_available(): - self.skipTest("Expert parallel mixin tests are CPU-only and should not run on GPU or XPU machines") + @is_tensor_parallel_test + def test_tp_backward(self): + """Tensor parallel backward (_tp_plan, no sequence parallel).""" + self._skip_if_mode_not_supported("tp") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - if os.cpu_count() < self.tensor_parallel_size: - self.skipTest( - f"Expert parallel tests require at least {self.tensor_parallel_size} CPUs, " - f"but only {os.cpu_count()} available" + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "tp" ) - if not hasattr(self.model_tester, "causal_lm_class") or self.model_tester.causal_lm_class is None: - self.skipTest("Model tester does not have causal_lm_class (not using CausalLMModelTester)") + @is_tensor_parallel_test + def test_sp_forward(self): + """Sequence-parallel forward (_sp_plan including per-layer MLP).""" + self._skip_if_mode_not_supported("sp") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - if not self._has_ep_plan(): - self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)(tmp_dir, model_class, atol, rtol, "sp") @is_tensor_parallel_test - def test_tp_forward(self): - self._skip_if_not_supported() - + def test_sp_backward(self): + """Sequence-parallel backward (_sp_plan including per-layer MLP).""" + self._skip_if_mode_not_supported("sp") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol with tempfile.TemporaryDirectory() as tmp_dir: set_seed(42) model = model_class(config) + self._assert_mixed_mlp_layers(model, config) model.save_pretrained(tmp_dir, save_original_format=True) - - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_forward_impl)(tmp_dir, model_class, atol, rtol) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "sp" + ) @is_tensor_parallel_test - def test_tp_backward(self): - self._skip_if_not_supported() - + def test_ep_forward(self): + """Expert-parallel forward (_ep_plan).""" + self._skip_if_mode_not_supported("ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol @@ -583,60 +567,146 @@ def test_tp_backward(self): set_seed(42) model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)(tmp_dir, model_class, atol, rtol, "ep") - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_backward_impl)(tmp_dir, model_class, atol, rtol) + @is_tensor_parallel_test + def test_ep_backward(self): + """Expert-parallel backward (_ep_plan).""" + self._skip_if_mode_not_supported("ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "ep" + ) @is_tensor_parallel_test def test_tp_generation(self): - # Test TP generation: unfused checkpoint → conversion mapping (if needed) → TP sharding → model → generate - self._skip_if_not_supported() - + """Tensor parallel generation (_tp_plan).""" + self._skip_if_mode_not_supported("tp") config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol - model_class = self._get_tp_model_class() + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_generation_impl)( + tmp_dir, model_class, atol, rtol, "tp", 25 + ) + + @is_tensor_parallel_test + def test_ep_generation(self): + """Expert-parallel generation (_ep_plan).""" + self._skip_if_mode_not_supported("ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol - max_new_tokens = 25 with tempfile.TemporaryDirectory() as tmp_dir: set_seed(42) model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) - _init_distributed(tp=self.tensor_parallel_size)(_test_tp_generation_impl)( - tmp_dir, model_class, atol, rtol, max_new_tokens + _init_distributed(tp=self.tensor_parallel_size)(_test_generation_impl)( + tmp_dir, model_class, atol, rtol, "ep", 25 ) @is_tensor_parallel_test def test_tp_generation_quantized(self): - self._skip_if_not_supported() + self._skip_if_mode_not_supported("tp") if not is_torchao_available(): self.skipTest("Test requires torchao") - self.skipTest("Quantization is not currently supported with distributed training") + self.skipTest("Quantization is not currently supported with distributed training (dtensor)") @is_tensor_parallel_test - def test_ep_forward(self): - self._skip_if_ep_not_supported() + def test_tp_ep_forward(self): + """Inference TP + EP forward (_tp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("tp_ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)( + tmp_dir, model_class, atol, rtol, "tp_ep" + ) + @is_tensor_parallel_test + def test_tp_ep_backward(self): + """Inference TP + EP backward (_tp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("tp_ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol with tempfile.TemporaryDirectory() as tmp_dir: set_seed(42) model = model_class(config) + self._assert_mixed_mlp_layers(model, config) model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "tp_ep" + ) - _init_distributed(tp=self.tensor_parallel_size)(_test_ep_forward_impl)(tmp_dir, model_class, atol, rtol) + @is_tensor_parallel_test + def test_sp_ep_forward(self): + """Training SP + EP forward (_sp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("sp_ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_forward_impl)( + tmp_dir, model_class, atol, rtol, "sp_ep" + ) @is_tensor_parallel_test - def test_ep_backward(self): - self._skip_if_ep_not_supported() + def test_sp_ep_backward(self): + """Training SP + EP backward (_sp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("sp_ep") + config = self.model_tester.get_config() + model_class = self._get_model_class() + atol = self.tensor_parallel_atol + rtol = self.tensor_parallel_rtol + + with tempfile.TemporaryDirectory() as tmp_dir: + set_seed(42) + model = model_class(config) + self._assert_mixed_mlp_layers(model, config) + model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_backward_impl)( + tmp_dir, model_class, atol, rtol, "sp_ep" + ) + @is_tensor_parallel_test + def test_tp_ep_generation(self): + """Inference TP + EP generation (_tp_plan ∪ _ep_plan).""" + self._skip_if_mode_not_supported("tp_ep") config = self.model_tester.get_config() - model_class = self._get_tp_model_class() + model_class = self._get_model_class() atol = self.tensor_parallel_atol rtol = self.tensor_parallel_rtol @@ -644,5 +714,15 @@ def test_ep_backward(self): set_seed(42) model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) + _init_distributed(tp=self.tensor_parallel_size)(_test_generation_impl)( + tmp_dir, model_class, atol, rtol, "tp_ep", 25 + ) + + @is_tensor_parallel_test + def test_tp_ep_generation_quantized(self): + self._skip_if_mode_not_supported("tp_ep") + + if not is_torchao_available(): + self.skipTest("Test requires torchao") - _init_distributed(tp=self.tensor_parallel_size)(_test_ep_backward_impl)(tmp_dir, model_class, atol, rtol) + self.skipTest("Quantization is not currently supported with distributed training (dtensor)")