From fc1076f348a49bce83f5f16db866e512a01f673f Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 7 Dec 2025 21:51:50 -0600 Subject: [PATCH 01/30] first attempt in supporting deepseek v3.2 --- .../megatron_to_hf/deepseekv32.py | 124 +++++++++++ miles/utils/data.py | 14 ++ miles/utils/deepseek_v32_patch.py | 50 +++++ miles_plugins/mbridge/__init__.py | 24 +- miles_plugins/mbridge/deepseekv32.py | 51 +++++ scripts/run_deepseek_v3.2_5layer.py | 206 ++++++++++++++++++ scripts/train_dsv32.py | 122 +++++++++++ 7 files changed, 590 insertions(+), 1 deletion(-) create mode 100644 miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py create mode 100644 miles/utils/deepseek_v32_patch.py create mode 100644 miles_plugins/mbridge/deepseekv32.py create mode 100644 scripts/run_deepseek_v3.2_5layer.py create mode 100644 scripts/train_dsv32.py diff --git a/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py b/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py new file mode 100644 index 00000000000..b271e5a1ed3 --- /dev/null +++ b/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py @@ -0,0 +1,124 @@ +import re + +import sglang +import torch +from packaging.version import parse + + +def convert_deepseekv3_2_to_hf(args, name, param): + if name == "module.module.embedding.word_embeddings.weight": + return [("model.embed_tokens.weight", param)] + if name == "module.module.output_layer.weight": + return [("lm_head.weight", param)] + if name == "module.module.decoder.final_layernorm.weight": + return [("model.norm.weight", param)] + + try: + head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads + except AttributeError: + head_dim = args.hidden_size // args.num_attention_heads + value_num_per_group = args.num_attention_heads // args.num_query_groups + + decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" + match = re.match(decoder_layers_pattern, name) + if match: + layer_idx, rest = match.groups() + + # experts + expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)" + match = re.match(expert_pattern, rest) + if match: + rest, expert_idx = match.groups() + if rest == "linear_fc1": + gate_weight, up_weight = param.chunk(2, dim=0) + outputs = [ + (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight), + (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight), + ] + return outputs + elif rest == "linear_fc2": + outputs = [ + (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param), + ] + if parse(sglang.__version__) < parse("0.4.9.post5") and args.sglang_enable_ep_moe: + outputs += [ + ( + f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.input_scale", + torch.tensor(1.0, dtype=torch.float32, device=param.device), + ), + ( + f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight_scale", + torch.tensor(1.0, dtype=torch.float32, device=param.device), + ), + ] + return outputs + else: + raise ValueError(f"Unknown expert parameter name: {name}") + + # shared expert + shared_expert_pattern = r"mlp.shared_experts\.(.+)" + match = re.match(shared_expert_pattern, rest) + if match: + rest = match.groups()[0] + if rest == "linear_fc1.weight": + gate_weight, up_weight = param.chunk(2, dim=0) + return [ + (f"model.layers.{layer_idx}.mlp.shared_experts.gate_proj.weight", gate_weight), + (f"model.layers.{layer_idx}.mlp.shared_experts.up_proj.weight", up_weight), + ] + elif rest == "linear_fc2.weight": + return [(f"model.layers.{layer_idx}.mlp.shared_experts.down_proj.weight", param)] + else: + raise ValueError(f"Unknown shared expert parameter name: {name}") + + if rest == "self_attention.linear_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)] + elif rest == "self_attention.linear_q_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_proj.weight", param)] + elif rest == "self_attention.linear_q_down_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_a_proj.weight", param)] + elif rest == "self_attention.linear_q_up_proj.layer_norm_weight": + return [(f"model.layers.{layer_idx}.self_attn.q_a_layernorm.weight", param)] + elif rest == "self_attention.linear_q_up_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.q_b_proj.weight", param)] + elif rest == "self_attention.linear_qkv.bias": + param = param.view(args.num_query_groups, -1) + q_bias, k_bias, v_bias = torch.split( + param, + split_size_or_sections=[value_num_per_group * head_dim, head_dim, head_dim], + dim=1, + ) + q_bias = q_bias.contiguous().flatten() + k_bias = k_bias.contiguous().flatten() + v_bias = v_bias.contiguous().flatten() + return [ + (f"model.layers.{layer_idx}.self_attn.q_proj.bias", q_bias), + (f"model.layers.{layer_idx}.self_attn.k_proj.bias", k_bias), + (f"model.layers.{layer_idx}.self_attn.v_proj.bias", v_bias), + ] + elif rest == "mlp.linear_fc1.weight": + gate_weight, up_weight = param.chunk(2, dim=0) + return [ + (f"model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight), + (f"model.layers.{layer_idx}.mlp.up_proj.weight", up_weight), + ] + elif rest == "mlp.linear_fc2.weight": + return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)] + elif rest == "self_attention.linear_qkv.layer_norm_weight" or rest == "input_layernorm.weight": + return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] + elif rest == "mlp.linear_fc1.layer_norm_weight": + return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] + elif rest == "self_attention.linear_kv_down_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.kv_a_proj_with_mqa.weight", param)] + elif rest == "self_attention.linear_kv_up_proj.layer_norm_weight": + return [(f"model.layers.{layer_idx}.self_attn.kv_a_layernorm.weight", param)] + elif rest == "self_attention.linear_kv_up_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.kv_b_proj.weight", param)] + elif rest == "pre_mlp_layernorm.weight": + return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] + elif rest == "mlp.router.weight": + return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)] + elif rest == "mlp.router.expert_bias": + return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", param)] + + raise ValueError(f"Unknown parameter name: {name}") diff --git a/miles/utils/data.py b/miles/utils/data.py index 6e64ef678de..737246acddd 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -207,6 +207,20 @@ def __init__( add_generation_prompt=True, **(apply_chat_template_kwargs or {}), ) + ### DSV32 + try: + prompt = tokenizer.apply_chat_template( + prompt, + tools, + tokenize=False, + add_generation_prompt=True, + **apply_chat_template_kwargs, + ) + except Exception as e: + from sglang.srt.entrypoints.openai.encoding_dsv32 import encode_messages + encode_config = dict(thinking_mode="thinking", drop_thinking=True, add_default_bos_token=True) + prompt = encode_messages(prompt, **encode_config) + ### DSV32 else: output_prompt = prompt diff --git a/miles/utils/deepseek_v32_patch.py b/miles/utils/deepseek_v32_patch.py new file mode 100644 index 00000000000..94009070650 --- /dev/null +++ b/miles/utils/deepseek_v32_patch.py @@ -0,0 +1,50 @@ +import os +import json +import tempfile +from transformers import AutoConfig + +_patched = False + + +def apply_deepseek_v32_patch(restore_model_type=False): + global _patched + + if _patched: + return + + _original_from_pretrained = AutoConfig.from_pretrained + + def _patched_from_pretrained(pretrained_model_name_or_path, *args, **kwargs): + if isinstance(pretrained_model_name_or_path, str) and os.path.isdir(pretrained_model_name_or_path): + config_file = os.path.join(pretrained_model_name_or_path, "config.json") + if os.path.exists(config_file): + try: + with open(config_file, "r") as f: + config_json = json.load(f) + + if config_json.get("model_type") == "deepseek_v32": + config_json["model_type"] = "deepseek_v3" + if "architectures" in config_json: + config_json["architectures"] = ["DeepseekV3ForCausalLM"] + + tmp_path = os.path.join(tempfile.gettempdir(), "_tmp_config_folder") + os.makedirs(tmp_path, exist_ok=True) + unique_path = os.path.join(tmp_path, f"deepseek_v32_{os.getpid()}.json") + + with open(unique_path, "w") as f: + json.dump(config_json, f) + + config = _original_from_pretrained(unique_path, *args, **kwargs) + + if restore_model_type: + object.__setattr__(config, "model_type", "deepseek_v32") + + return config + except Exception: + pass + + return _original_from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + + AutoConfig.from_pretrained = _patched_from_pretrained + _patched = True + diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index f97c7f46eef..0e259d6b0f5 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -1,6 +1,28 @@ +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) + +from miles.utils.deepseek_v32_patch import apply_deepseek_v32_patch +apply_deepseek_v32_patch(restore_model_type=True) + +from .deepseekv32 import DeepseekV32Bridge from .glm4 import GLM4Bridge from .glm4moe import GLM4MoEBridge from .mimo import MimoBridge from .qwen3_next import Qwen3NextBridge -__all__ = ["GLM4Bridge", "GLM4MoEBridge", "Qwen3NextBridge", "MimoBridge"] +__all__ = ["DeepseekV32Bridge", "GLM4Bridge", "GLM4MoEBridge", "Qwen3NextBridge", "MimoBridge"] + +from mbridge import AutoBridge + +_original_from_config = AutoBridge.from_config + +@classmethod +def _patched_from_config(cls, hf_config, **kwargs): + if hf_config.model_type == "deepseek_v32": + from mbridge.core.bridge import _MODEL_REGISTRY + return _MODEL_REGISTRY['deepseek_v32'](hf_config, **kwargs) + + return _original_from_config(hf_config, **kwargs) + +AutoBridge.from_config = _patched_from_config diff --git a/miles_plugins/mbridge/deepseekv32.py b/miles_plugins/mbridge/deepseekv32.py new file mode 100644 index 00000000000..8bc94dd23ff --- /dev/null +++ b/miles_plugins/mbridge/deepseekv32.py @@ -0,0 +1,51 @@ +from mbridge.core import register_model +from mbridge.models import DeepseekV3Bridge + + +@register_model("deepseek_v32") +class DeepseekV32Bridge(DeepseekV3Bridge): + + _ATTENTION_MAPPING = ( + DeepseekV3Bridge._ATTENTION_MAPPING.copy() + ) + + # Because the indexer needs the norm output, we cannot use the fused transformer engine impl and have to compute it separately. + if "self_attention.linear_q_up_proj.layer_norm_weight" in _ATTENTION_MAPPING: + del _ATTENTION_MAPPING["self_attention.linear_q_up_proj.layer_norm_weight"] + if "self_attention.linear_kv_up_proj.layer_norm_weight" in _ATTENTION_MAPPING: + del _ATTENTION_MAPPING["self_attention.linear_kv_up_proj.layer_norm_weight"] + + _ATTENTION_MAPPING.update({ + "self_attention.q_layernorm.weight": [ + "model.layers.{layer_number}.self_attn.q_a_layernorm.weight" + ], + "self_attention.kv_layernorm.weight": [ + "model.layers.{layer_number}.self_attn.kv_a_layernorm.weight" + ], + "self_attention.core_attention.indexer.linear_wq_b.weight": [ + "model.layers.{layer_number}.self_attn.indexer.wq_b.weight" + ], + "self_attention.core_attention.indexer.linear_wk.weight": [ + "model.layers.{layer_number}.self_attn.indexer.wk.weight" + ], + "self_attention.core_attention.indexer.k_norm.weight": [ + "model.layers.{layer_number}.self_attn.indexer.k_norm.weight" + ], + "self_attention.core_attention.indexer.k_norm.bias": [ + "model.layers.{layer_number}.self_attn.indexer.k_norm.bias" + ], + "self_attention.core_attention.indexer.linear_weights_proj.weight": [ + "model.layers.{layer_number}.self_attn.indexer.weights_proj.weight" + ], + }) + + def _build_config(self): + config = super()._build_config() + + config.experimental_attention_variant = "dsa" + config.dsa_indexer_n_heads = getattr(self.hf_config, 'dsa_indexer_n_heads', 64) + config.dsa_indexer_head_dim = getattr(self.hf_config, 'dsa_indexer_head_dim', 128) + config.dsa_indexer_topk = getattr(self.hf_config, 'dsa_indexer_topk', 2048) + + return config + diff --git a/scripts/run_deepseek_v3.2_5layer.py b/scripts/run_deepseek_v3.2_5layer.py new file mode 100644 index 00000000000..0923e11c68f --- /dev/null +++ b/scripts/run_deepseek_v3.2_5layer.py @@ -0,0 +1,206 @@ +import re +from dataclasses import dataclass +from typing import Literal + +import typer + +import miles.utils.external_utils.command_utils as U + +app = typer.Typer() + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + run_id: str = U.create_run_id() + hf_checkpoint: str = "/root/.cache/dsv32-ckpt/DeepSeek-V3.2-5layer" + torch_dist_checkpoint: str = "/root/.cache/dsv32-ckpt/DeepSeek-V3-0324-5layer_torch_dist" + num_gpus_per_node: int = 8 + enable_eval: bool = False + enable_deepep: bool = False + extra_args: str = "" + task: Literal["dapo_aime", "gsm8k"] = "dapo_aime" + mode: Literal["normal", "debug_minimal"] = "debug_minimal" + + +@app.command() +@U.dataclass_cli +def train(args: ScriptArgs): + load_save_path = f"/root/shared_data/{args.run_id}/checkpoints" + ckpt_args = ( + f"--hf-checkpoint {args.hf_checkpoint} " + f"--ref-load {args.torch_dist_checkpoint} " + f"--load {load_save_path} " + f"--save {load_save_path} " + "--save-interval 20 " + "--save-retain-interval 20 " + ) + + rollout_args = ( + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 3000 " + "--rollout-batch-size 128 " + "--n-samples-per-prompt 8 " + "--rollout-temperature 0.8 " + "--num-steps-per-rollout 4 " + "--balance-data " + ) + + if args.mode != "debug_minimal": + rollout_args += ( + "--over-sampling-batch-size 256 " + "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " + ) + + eval_args = "" + if (args.mode != "debug_minimal") and args.enable_eval: + eval_args += "--eval-interval 20 " "--eval-top-p 0.7 " + + match args.task: + case "dapo_aime": + rollout_args += ( + "--prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 32768} " + ) + eval_args += ( + "--eval-prompt-data aime /root/aime-2024/aime-2024.jsonl " + "--n-samples-per-eval-prompt 8 " + "--eval-max-response-len 32768 " + ) + case "gsm8k": + rollout_args += ( + "--prompt-data /root/gsm8k/train.parquet " + "--input-key messages " + "--rollout-max-response-len 256 " + ) + eval_args += ( + "--eval-prompt-data gsm8k /root/gsm8k/test.parquet " + "--n-samples-per-eval-prompt 1 " + "--eval-max-response-len 256 " + ) + + if args.num_nodes <= 2: + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 8 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + ) + elif args.num_nodes <= 4: + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 8 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + ) + else: + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 8 " + "--expert-model-parallel-size 16 " + "--expert-tensor-parallel-size 1 " + ) + perf_args += ( + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 2048 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_decode_max_bs = 256 + sglang_world_size = 8 if args.num_nodes <= 4 else 64 + sglang_attn_dp_size = 1 if args.num_nodes <= 4 else 8 + sglang_attn_tp_size = sglang_world_size // sglang_attn_dp_size + sglang_args = ( + f"--rollout-num-gpus-per-engine {sglang_world_size} " + "--sglang-mem-fraction-static 0.7 " + # f"--sglang-tp-size {sglang_world_size} " + f"--sglang-tp-size 1 " + f"--sglang-ep-size {sglang_world_size} " + "--sglang-enable-dp-attention " + f"--sglang-dp-size {sglang_attn_dp_size} " + "--sglang-moe-dense-tp-size 1 " + "--sglang-enable-dp-lm-head " + "--sglang-server-concurrency 1024 " + f"--sglang-max-running-requests {sglang_world_size * sglang_decode_max_bs // sglang_attn_tp_size} " + f"--sglang-chunked-prefill-size {sglang_world_size * sglang_decode_max_bs} " + f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} " + ) + if args.enable_deepep: + sglang_args += ( + "--sglang-moe-a2a-backend deepep " + "--sglang-deepep-mode low_latency " + ) + sglang_extra_env_vars = {} + if args.enable_deepep: + sglang_extra_env_vars["SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = f"{sglang_decode_max_bs}" + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + f"--update-weight-buffer-size {4 * 1024 ** 3} " + f"--actor-num-nodes {args.num_nodes} " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + f"--num-gpus-per-node {args.num_gpus_per_node} " + "--colocate " + "--use-fault-tolerance " + f"--dump-details /root/shared_data/{args.run_id}/dump_details " + "--disable-weights-backuper " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{misc_args} " + f"{args.extra_args} " + ) + + U.execute_train( + train_args=train_args, + train_script="scripts/train_dsv32.py", + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type="deepseek-v32-5layer", + extra_env_vars={**sglang_extra_env_vars}, + ) + + +if __name__ == "__main__": + app() + diff --git a/scripts/train_dsv32.py b/scripts/train_dsv32.py new file mode 100644 index 00000000000..4216ec4bec6 --- /dev/null +++ b/scripts/train_dsv32.py @@ -0,0 +1,122 @@ +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from miles.utils.deepseek_v32_patch import apply_deepseek_v32_patch +apply_deepseek_v32_patch() + +from turtle import mode +import ray +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS +from typing import Optional + +try: + from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH +except ImportError: + GPU_MEMORY_TYPE_CUDA_GRAPH = None + +from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models +from miles.utils.arguments import parse_args +from miles.utils.logging_utils import configure_logger +from miles.utils.tracking_utils import init_tracking + + +def train(args): + configure_logger() + # allocate the GPUs + pgs = create_placement_groups(args) + init_tracking(args) + + # create the rollout manager, with sglang engines inside. + # need to initialize rollout manager first to calculate num_rollout + rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + + # create the actor and critic models + actor_model, critic_model = create_training_models(args, pgs, rollout_manager) + + if args.offload_rollout: + ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])) + + # always update weight first so that sglang has the loaded weights from training. + print("[DEBUG] train.py first update weights") + actor_model.update_weights() + + if args.check_weight_update_equal: + ray.get(rollout_manager.check_weights.remote(action="compare")) + + if args.offload_rollout: + if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: + ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])) + ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE])) + + # special case for eval-only + if args.num_rollout == 0 and args.eval_interval is not None: + ray.get(rollout_manager.eval.remote(rollout_id=0)) + + def offload_train(): + if args.offload_train: + if args.use_critic: + critic_model.offload() + if rollout_id >= args.num_critic_only_steps: + actor_model.offload() + else: + actor_model.offload() + else: + actor_model.clear_memory() + + def onload_rollout(): + if args.offload_rollout: + ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])) + + # train loop. + # note that for async training, one can change the position of the sync operation(ray.get). + for rollout_id in range(args.start_rollout_id, args.num_rollout): + # TODO extract the duplicated eval logic + if args.eval_interval is not None and rollout_id == 0: + ray.get(rollout_manager.eval.remote(rollout_id)) + + rollout_data_ref = ray.get(rollout_manager.generate.remote(rollout_id)) + + if args.offload_rollout: + ray.get(rollout_manager.offload.remote()) + + if args.use_critic: + critic_train_handle = critic_model.async_train(rollout_id, rollout_data_ref) + if rollout_id >= args.num_critic_only_steps: + ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) + ray.get(critic_train_handle) + else: + ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) + + if args.save_interval is not None and ( + (rollout_id + 1) % args.save_interval == 0 + or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0) + ): + if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps): + actor_model.save_model(rollout_id) + if args.use_critic: + critic_model.save_model(rollout_id) + if args.rollout_global_dataset: + ray.get(rollout_manager.save.remote(rollout_id)) + + offload_train() + onload_rollout() + actor_model.update_weights() + + if args.offload_rollout: + if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: + ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])) + ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE])) + + if args.eval_interval is not None and ( + (rollout_id + 1) % args.eval_interval == 0 + or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0) + ): + ray.get(rollout_manager.eval.remote(rollout_id)) + + ray.get(rollout_manager.dispose.remote()) + + +if __name__ == "__main__": + args = parse_args() + train(args) From a7373e7a441dd47abb8c869ed77ae74ace052dc9 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 10 Dec 2025 11:32:39 -0800 Subject: [PATCH 02/30] update --- miles/backends/megatron_utils/actor.py | 3 +++ .../megatron_utils/megatron_to_hf/__init__.py | 3 +++ .../megatron_to_hf/deepseekv32.py | 17 ++++++++++++++--- miles_plugins/mbridge/deepseekv32.py | 3 +++ scripts/run_deepseek_v3.2_5layer.py | 12 +++++++----- scripts/train_dsv32.py | 1 - 6 files changed, 30 insertions(+), 9 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index a92198a6744..a12d1be4c95 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -13,6 +13,9 @@ from torch_memory_saver import torch_memory_saver from transformers import AutoConfig, AutoTokenizer +from miles.utils.deepseek_v32_patch import apply_deepseek_v32_patch +apply_deepseek_v32_patch() + from miles.ray.train_actor import TrainRayActor from miles.utils import train_dump_utils from miles.utils.context_utils import with_defer diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py index 84ff899aa52..b9b394cbcff 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py @@ -1,4 +1,5 @@ from .deepseekv3 import convert_deepseekv3_to_hf +from .deepseekv32 import convert_deepseekv32_to_hf from .glm4 import convert_glm4_to_hf from .glm4moe import convert_glm4moe_to_hf from .llama import convert_llama_to_hf @@ -41,6 +42,8 @@ def _convert_to_hf_core(args, model_name, name, param): converted_named_tensors = convert_qwen3_next_to_hf(args, name, param) elif "qwen2" in model_name or "qwen3" in model_name: converted_named_tensors = convert_qwen2_to_hf(args, name, param) + elif "deepseekv32" in model_name: + converted_named_tensors = convert_deepseekv32_to_hf(args, name, param) elif "deepseekv3" in model_name: converted_named_tensors = convert_deepseekv3_to_hf(args, name, param) diff --git a/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py b/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py index b271e5a1ed3..3b6519ada8e 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py +++ b/miles/backends/megatron_utils/megatron_to_hf/deepseekv32.py @@ -5,7 +5,7 @@ from packaging.version import parse -def convert_deepseekv3_2_to_hf(args, name, param): +def convert_deepseekv32_to_hf(args, name, param): if name == "module.module.embedding.word_embeddings.weight": return [("model.embed_tokens.weight", param)] if name == "module.module.output_layer.weight": @@ -77,7 +77,7 @@ def convert_deepseekv3_2_to_hf(args, name, param): return [(f"model.layers.{layer_idx}.self_attn.q_proj.weight", param)] elif rest == "self_attention.linear_q_down_proj.weight": return [(f"model.layers.{layer_idx}.self_attn.q_a_proj.weight", param)] - elif rest == "self_attention.linear_q_up_proj.layer_norm_weight": + elif rest == "self_attention.q_layernorm.weight": return [(f"model.layers.{layer_idx}.self_attn.q_a_layernorm.weight", param)] elif rest == "self_attention.linear_q_up_proj.weight": return [(f"model.layers.{layer_idx}.self_attn.q_b_proj.weight", param)] @@ -110,12 +110,23 @@ def convert_deepseekv3_2_to_hf(args, name, param): return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] elif rest == "self_attention.linear_kv_down_proj.weight": return [(f"model.layers.{layer_idx}.self_attn.kv_a_proj_with_mqa.weight", param)] - elif rest == "self_attention.linear_kv_up_proj.layer_norm_weight": + elif rest == "self_attention.kv_layernorm.weight": return [(f"model.layers.{layer_idx}.self_attn.kv_a_layernorm.weight", param)] elif rest == "self_attention.linear_kv_up_proj.weight": return [(f"model.layers.{layer_idx}.self_attn.kv_b_proj.weight", param)] elif rest == "pre_mlp_layernorm.weight": return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] + # DSA Indexer parameters + elif rest == "self_attention.core_attention.indexer.linear_wq_b.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.wq_b.weight", param)] + elif rest == "self_attention.core_attention.indexer.linear_wk.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.wk.weight", param)] + elif rest == "self_attention.core_attention.indexer.k_norm.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.k_norm.weight", param)] + elif rest == "self_attention.core_attention.indexer.k_norm.bias": + return [(f"model.layers.{layer_idx}.self_attn.indexer.k_norm.bias", param)] + elif rest == "self_attention.core_attention.indexer.linear_weights_proj.weight": + return [(f"model.layers.{layer_idx}.self_attn.indexer.weights_proj.weight", param)] elif rest == "mlp.router.weight": return [(f"model.layers.{layer_idx}.mlp.gate.weight", param)] elif rest == "mlp.router.expert_bias": diff --git a/miles_plugins/mbridge/deepseekv32.py b/miles_plugins/mbridge/deepseekv32.py index 8bc94dd23ff..fb9355f5ee3 100644 --- a/miles_plugins/mbridge/deepseekv32.py +++ b/miles_plugins/mbridge/deepseekv32.py @@ -1,5 +1,6 @@ from mbridge.core import register_model from mbridge.models import DeepseekV3Bridge +from megatron.core.transformer.enums import AttnBackend @register_model("deepseek_v32") @@ -42,6 +43,8 @@ class DeepseekV32Bridge(DeepseekV3Bridge): def _build_config(self): config = super()._build_config() + config.attention_backend = AttnBackend.auto + config.experimental_attention_variant = "dsa" config.dsa_indexer_n_heads = getattr(self.hf_config, 'dsa_indexer_n_heads', 64) config.dsa_indexer_head_dim = getattr(self.hf_config, 'dsa_indexer_head_dim', 128) diff --git a/scripts/run_deepseek_v3.2_5layer.py b/scripts/run_deepseek_v3.2_5layer.py index 0923e11c68f..bcbb3a891af 100644 --- a/scripts/run_deepseek_v3.2_5layer.py +++ b/scripts/run_deepseek_v3.2_5layer.py @@ -13,8 +13,8 @@ class ScriptArgs(U.ExecuteTrainConfig): run_id: str = U.create_run_id() hf_checkpoint: str = "/root/.cache/dsv32-ckpt/DeepSeek-V3.2-5layer" - torch_dist_checkpoint: str = "/root/.cache/dsv32-ckpt/DeepSeek-V3-0324-5layer_torch_dist" - num_gpus_per_node: int = 8 + torch_dist_checkpoint: str = "/root/DeepSeek-V3.2-5layer_torch_dist" + num_gpus_per_node: int = 4 enable_eval: bool = False enable_deepep: bool = False extra_args: str = "" @@ -87,8 +87,8 @@ def train(args: ScriptArgs): "--tensor-model-parallel-size 1 " "--sequence-parallel " "--pipeline-model-parallel-size 1 " - "--context-parallel-size 8 " - "--expert-model-parallel-size 8 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 4 " "--expert-tensor-parallel-size 1 " ) elif args.num_nodes <= 4: @@ -136,7 +136,7 @@ def train(args: ScriptArgs): ) sglang_decode_max_bs = 256 - sglang_world_size = 8 if args.num_nodes <= 4 else 64 + sglang_world_size = 4 if args.num_nodes <= 4 else 64 sglang_attn_dp_size = 1 if args.num_nodes <= 4 else 8 sglang_attn_tp_size = sglang_world_size // sglang_attn_dp_size sglang_args = ( @@ -168,6 +168,7 @@ def train(args: ScriptArgs): "--hidden-dropout 0.0 " "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " + "--attention-backend auto " f"--update-weight-buffer-size {4 * 1024 ** 3} " f"--actor-num-nodes {args.num_nodes} " f"--actor-num-gpus-per-node {args.num_gpus_per_node} " @@ -176,6 +177,7 @@ def train(args: ScriptArgs): "--use-fault-tolerance " f"--dump-details /root/shared_data/{args.run_id}/dump_details " "--disable-weights-backuper " + "--model-name deepseekv32 " ) train_args = ( diff --git a/scripts/train_dsv32.py b/scripts/train_dsv32.py index 4216ec4bec6..848cef803fb 100644 --- a/scripts/train_dsv32.py +++ b/scripts/train_dsv32.py @@ -38,7 +38,6 @@ def train(args): ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])) # always update weight first so that sglang has the loaded weights from training. - print("[DEBUG] train.py first update weights") actor_model.update_weights() if args.check_weight_update_equal: From b62966e0eb7fa1e404051b2444b5aa9d4d7f4e4b Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sat, 13 Dec 2025 19:06:18 -0800 Subject: [PATCH 03/30] add several fix, supported thd + CP on megatron's dsa, added dockerfile --- docker/deepseekv32/Dockerfile | 129 ++++ docker/deepseekv32/megatron_dsv32.patch | 552 ++++++++++++++++++ .../processors/quantizer_fp8.py | 7 +- miles/backends/training_utils/loss.py | 2 +- 4 files changed, 686 insertions(+), 4 deletions(-) create mode 100644 docker/deepseekv32/Dockerfile create mode 100644 docker/deepseekv32/megatron_dsv32.patch diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile new file mode 100644 index 00000000000..5787b1ce478 --- /dev/null +++ b/docker/deepseekv32/Dockerfile @@ -0,0 +1,129 @@ +ARG SGLANG_IMAGE_TAG=dev +FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang + +# ======================================== Arguments ============================================= + +ARG PATCH_VERSION=latest +ARG MEGATRON_COMMIT=436065a86b749ca3b50eebca68f55c9e690a9f63 + +ARG ENABLE_CUDA_13=0 + +ARG ENABLE_SGLANG_PATCH=0 + +# ======================================== Setup ============================================= + +WORKDIR /root/ + +# ======================================== Apt dependencies ============================================= + +RUN apt update +RUN apt install -y nvtop rsync dnsutils + +# ====================================== Python dependencies ============================================ + +# The compilation is slow, thus should be put at top +# TransformerEngines does not support too high FA2 +RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation + +# The compilation is slow, thus should be put at top +RUN git clone https://github.com/Dao-AILab/flash-attention.git && \ + cd flash-attention/ && git checkout fbf24f67cf7f6442c5cfb2c1057f4bfc57e72d89 && git submodule update --init && cd hopper/ && \ + MAX_JOBS=96 python setup.py install && \ + export python_path=`python -c "import site; print(site.getsitepackages()[0])"` && \ + mkdir -p $python_path/flash_attn_3 && \ + cp flash_attn_interface.py $python_path/flash_attn_3/flash_attn_interface.py && \ + rm -rf flash-attention/ + +RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps + +RUN pip install flash-linear-attention==0.4.0 + +# TE does not have wheel on cuda 13 yet, thus need to install from source +RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ + pip install nvidia-mathdx==25.6.0 && \ + pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.8; \ + else \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.8.0"; \ + fi + +RUN NVCC_APPEND_FLAGS="--threads 4" \ + pip -v install --disable-pip-version-check --no-cache-dir \ + --no-build-isolation \ + --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4 + +RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ + cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \ + pip install -e . + +RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall +RUN pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation +RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation + +# This patch from masahi will be included in later Triton releases +RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ + (cd /root && git clone -b feat/v350_plus_8045 https://github.com/fzyzcjy/triton.git && cd triton && pip install -r python/requirements.txt && pip install --verbose -e .); \ + fi + +COPY requirements.txt /tmp/requirements.txt +RUN pip install -r /tmp/requirements.txt + +# Temporarily install another sgl-kernel version for GB300 without rebuilding the whole image +RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ + SGL_KERNEL_VERSION=0.3.17.post2 && \ + python3 -m pip install https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu130-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps; \ + fi + +# This patch is merged into main, but we are using stable version, thus still need it +RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ + curl -L https://github.com/NVIDIA/TransformerEngine/pull/2286.patch -o /root/te2286.patch && (cd /usr/local/lib/python3.12/dist-packages/transformer_engine && (patch -p2 < /root/te2286.patch)); \ + fi + +# AMEM +# we need to create a fake libcuda.so.1 to make the linker happy when building AMEM +ENV CUDA_DIR=/usr/local/cuda +ENV CUDA_STUBS=${CUDA_DIR}/lib64/stubs +RUN ln -s ${CUDA_STUBS}/libcuda.so ${CUDA_STUBS}/libcuda.so.1 && \ + echo "${CUDA_STUBS}" > /etc/ld.so.conf.d/z-cuda-stubs.conf && \ + ldconfig +RUN git clone https://github.com/inclusionAI/asystem-amem.git && \ + cd asystem-amem && git checkout 6483bb17c9a98b51c3a94b7048467d5b50fbad4b && \ + git submodule init && git submodule update && \ + MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi/ ./build.sh && \ + mv /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2 /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2.bak && \ + cp -r third_party/nccl/build/lib/* /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/ + +RUN [ ! -f /root/.tmux.conf ] || rm /root/.tmux.conf + +# ====================================== Patches ============================================ + +COPY docker/deepseekv32/${PATCH_VERSION}/megatron_dsv32.patch /root/Megatron-LM/ +RUN cd Megatron-LM && \ + git update-index --refresh && \ + git apply megatron.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm megatron.patch + +# TODO temporarily skip patching for GB200/GB300 (and require users to bring their own sglang version). should add back later. +COPY docker/patch/${PATCH_VERSION}/sglang.patch /sgl-workspace/sglang/ +RUN if [ "$ENABLE_SGLANG_PATCH" = "1" ]; then \ + cd /sgl-workspace/sglang && \ + git update-index --refresh && \ + git apply sglang.patch && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm sglang.patch; \ +fi + +# ====================================== Install main package ============================================ + +# TODO may improve +ARG MILES_COMMIT=main +RUN git clone https://github.com/radixark/miles.git /root/miles && \ + cd /root/miles && \ + git checkout ${MILES_COMMIT} && \ + pip install -e . --no-deps diff --git a/docker/deepseekv32/megatron_dsv32.patch b/docker/deepseekv32/megatron_dsv32.patch new file mode 100644 index 00000000000..5ad0d3afe08 --- /dev/null +++ b/docker/deepseekv32/megatron_dsv32.patch @@ -0,0 +1,552 @@ +diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py +index 89659a1d7..f1d6855ee 100644 +--- a/megatron/core/transformer/dot_product_attention_context_parallel.py ++++ b/megatron/core/transformer/dot_product_attention_context_parallel.py +@@ -132,10 +132,10 @@ class AllGatherComm: + self.handles = [] + + +-def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stride, device, dtype): ++def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stride, device, dtype, if_zz_mask=False): + '''Convert the attention mask to the attention bias''' + +- if cp_size == 1: ++ if cp_size == 1 or if_zz_mask: + zz_mask = attention_mask + else: + chunked = attention_mask.chunk(dim=3, chunks=cp_size * 2) +@@ -143,7 +143,7 @@ def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stri + zz_mask = torch.cat(zz_mask, dim=3) + attn_bias = torch.zeros(zz_mask.shape, device=device, dtype=dtype) + attn_bias.masked_fill_(zz_mask, float('-inf')) +- attn_bias = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) ++ attn_bia = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) + return attn_bias + + +@@ -151,7 +151,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + """Native attention function with context parallelism.""" + + @staticmethod +- def forward(ctx, q, k, v, attention_mask, attention_dropout, softmax_scale, pg): ++ def forward(ctx, q, k, v, attention_mask, attention_dropout, softmax_scale, pg, if_zz_mask=False): + '''Forward pass for the native attention function with context parallelism''' + + # Assert einops exists +@@ -171,12 +171,17 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + probs = [] + + # Initialize KV buffers +- kv_buffer = torch.empty( +- (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), ++ # seperate KV buffer for MLA ++ kv_buffer = [torch.empty( ++ (k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), + dtype=k.dtype, + device=k.device, +- ) +- kv_buffer_copy = torch.empty_like(kv_buffer) ++ ), torch.empty( ++ (v.shape[0] * cp_size, v.shape[1], heads_k_stride, v.shape[3]), ++ dtype=v.dtype, ++ device=v.device, ++ )] ++ kv_buffer_copy = [torch.empty_like(kv_buffer[0]), torch.empty_like(kv_buffer[1])] + + # All-gather first chunk of KV buffers + k_0 = k[:, :, :heads_k_stride].contiguous() +@@ -186,7 +191,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + + # Prepare attention bias + attn_bias = to_zz_mask_attn_bias( +- attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype ++ attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype, if_zz_mask + ) + + # Iterate over heads +@@ -226,6 +231,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + + # Save contexts for backward pass + ctx.save_for_backward(q, k, v, attention_mask, *outs, *probs) ++ ctx.if_zz_mask = if_zz_mask + ctx.dropout = attention_dropout + ctx.scale = softmax_scale + ctx.heads_k_stride = heads_k_stride # TODO make it configurable +@@ -252,12 +258,16 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + comm = AllGatherComm(group=pg) + + # Initialize KV buffers +- kv_buffer = torch.empty( +- (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), ++ kv_buffer = [torch.empty( ++ (k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), + dtype=k.dtype, + device=k.device, +- ) +- kv_buffer_copy = torch.empty_like(kv_buffer) ++ ), torch.empty( ++ (v.shape[0] * cp_size, v.shape[1], heads_k_stride, v.shape[3]), ++ dtype=v.dtype, ++ device=v.device, ++ )] ++ kv_buffer_copy = [torch.empty_like(kv_buffer[0]), torch.empty_like(kv_buffer[1])] + + # All-gather first chunk of KV buffers + dq = [] +@@ -270,7 +280,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + + # Prepare attention bias + attn_bias = to_zz_mask_attn_bias( +- attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype ++ attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype, ctx.if_zz_mask + ) + + # Iterate over heads +@@ -339,4 +349,4 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + dq = torch.cat(dq, dim=2) + dk = torch.cat(dk, dim=2) + dv = torch.cat(dv, dim=2) +- return dq, dk, dv, None, None, None, None ++ return dq, dk, dv, None, None, None, None, None +diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py +index fc994490b..7bc9a485e 100644 +--- a/megatron/core/transformer/experimental_attention_variant/dsa.py ++++ b/megatron/core/transformer/experimental_attention_variant/dsa.py +@@ -6,6 +6,7 @@ from dataclasses import dataclass + from typing import Optional, Tuple, Union + + import torch ++import einops + + from megatron.core import parallel_state + from megatron.core.models.common.embeddings import ( +@@ -21,6 +22,8 @@ from megatron.core.transformer.module import MegatronModule + from megatron.core.transformer.spec_utils import ModuleSpec, build_module + from megatron.core.transformer.transformer_config import TransformerConfig + ++from megatron.core.transformer.dot_product_attention_context_parallel import AllGatherComm, AttentionFuncionWithContextParallel ++ + try: + from fast_hadamard_transform import hadamard_transform + except ImportError: +@@ -191,44 +194,72 @@ def compute_dsa_indexer_loss( + Returns: + index_loss: KL divergence loss (scalar). + """ +- sq, b, np, hn = query.size() +- sk = key.size(0) ++ cp_size = parallel_state.get_context_parallel_world_size() + +- # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] +- query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) +- # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] +- key = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) +- # Compute attention scores [b * np, sq, sk] +- attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale +- # Reshape to [b, np, sq, sk] +- attention_scores = attention_scores.reshape(b, np, sq, sk) ++ if cp_size > 1: ++ sq_local, b, np, hn = query.size() ++ sk_local = key.size(0) ++ sk_global = sk_local * cp_size + +- # causal_mask [sq, sk] +- causal_mask = torch.triu( +- torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), +- diagonal=1, +- ) +- # index_mask [b, sq, sk] +- index_mask = torch.full( +- (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device +- ).scatter_(-1, topk_indices, 0) +- +- # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] +- attention_scores += causal_mask.view(1, 1, sq, sk) +- if sparse_loss: +- # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] +- attention_scores += index_mask.view(b, 1, sq, sk) +- # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] +- index_scores += index_mask +- +- # [b, np, sq, sk] -> [b, np, sq, sk] +- attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) +- # [b, sq, sk] -> [b, sq, sk] +- index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) ++ causal_mask = get_causal_mask(sq_local, sk_local, query.device) ++ float_mask = torch.zeros_like(causal_mask, dtype=torch.float32).masked_fill( ++ causal_mask, float('-inf') ++ ) + +- # Sum attention scores across heads. +- # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] +- attention_scores = attention_scores.sum(dim=1) ++ index_mask = torch.full( ++ (b, sq_local, sk_global), float("-inf"), dtype=torch.float32, device=causal_mask.device ++ ).scatter_(-1, topk_indices, 0) ++ ++ float_mask = float_mask.view(1, 1, sq_local, sk_global) ++ float_mask = index_mask.view(b, 1, sq_local, sk_global) + float_mask if sparse_loss else float_mask ++ ++ # because the attention computation is more heavy in memory (has head dim), ++ # we apply cp (all-gather backend) on attention scores computation ++ attention_scores = compute_attention_scores_with_cp(query, key, float_mask, softmax_scale) # [b, sq_local, sk_global] ++ ++ index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) ++ ++ else: ++ sq, b, np, hn = query.size() ++ sk = key.size(0) ++ ++ # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] ++ query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) ++ # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] ++ key = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) ++ # Compute attention scores [b * np, sq, sk] ++ attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale ++ # Reshape to [b, np, sq, sk] ++ attention_scores = attention_scores.reshape(b, np, sq, sk) ++ ++ # causal_mask [sq, sk] ++ causal_mask = torch.triu( ++ torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), ++ diagonal=1, ++ ) ++ # index_mask [b, sq, sk] ++ index_mask = torch.full( ++ (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device ++ ).scatter_(-1, topk_indices, 0) ++ ++ # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] ++ attention_scores += causal_mask.view(1, 1, sq, sk) ++ if sparse_loss: ++ # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] ++ attention_scores += index_mask.view(b, 1, sq, sk) ++ # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] ++ index_scores += index_mask ++ ++ # [b, np, sq, sk] -> [b, np, sq, sk] ++ attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) ++ # [b, sq, sk] -> [b, sq, sk] ++ index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) ++ ++ # Sum attention scores across heads. ++ # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] ++ attention_scores = attention_scores.sum(dim=1) ++ ++ # Common part + if pg_collection.tp.size() > 1: + # attention scores are scattered to TP ranks in head dimension. + torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) +@@ -252,6 +283,57 @@ def compute_dsa_indexer_loss( + return indexer_loss + + ++def compute_attention_scores_with_cp(q, k, attn_bias, scale, heads_k_stride = 1): ++ """ ++ compute attention scores of q_local @ k_global with CP all-gather backend ++ parallel on n_heads dimension ++ """ ++ pg = parallel_state.get_context_parallel_group() ++ cp_size = parallel_state.get_context_parallel_world_size() ++ ++ sq_local, b, nheads, hn_q = q.shape ++ sk_local, _, nheads_k, hn_k = k.shape ++ sk_global = sk_local * cp_size ++ ++ assert nheads % nheads_k == 0 and nheads_k % heads_k_stride == 0 ++ ++ comm = AllGatherComm(group=pg) ++ attns = torch.zeros(b, heads_k_stride, sq_local, sk_global, dtype=q.dtype, device=q.device) ++ ++ k_buffer = torch.empty( ++ (sk_global, b, heads_k_stride, hn_k), ++ dtype=k.dtype, ++ device=k.device ++ ) ++ k_buffer_copy = torch.empty_like(k_buffer) ++ k_0 = k[:, :, :heads_k_stride].contiguous() ++ comm.all_gather(k_buffer_copy, k_0) ++ ++ attn_bias = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) ++ ++ for i in range(0, nheads_k, heads_k_stride): ++ comm.wait() ++ k_buffer, k_buffer_copy = k_buffer_copy, k_buffer ++ if i < nheads_k - heads_k_stride: ++ kvsl = i + heads_k_stride ++ kvsr = kvsl + heads_k_stride ++ send_k = k[:, :, kvsl:kvsr].contiguous() ++ comm.all_gather(k_buffer_copy, send_k) ++ q_i = q[:, :, i * nheads // nheads_k : (i + heads_k_stride) * nheads // nheads_k] ++ k_i = k_buffer ++ ++ _q_i = einops.rearrange(q_i, 's b h d -> b h s d') ++ _k_i = einops.rearrange(k_i, 's b h d -> b h d s') ++ attn_i = torch.matmul(_q_i.float(), _k_i.float()) * scale + attn_bias ++ attn_i = torch.nn.functional.softmax(attn_i, dim=-1, dtype=torch.float32) ++ ++ attns = attns + attn_i ++ ++ attns = torch.sum(attns, dim=1) ++ ++ return attns ++ ++ + class DSAIndexerLossAutoScaler(torch.autograd.Function): + """An AutoScaler that triggers the backward pass and scales the grad for indexer loss. + +@@ -496,7 +578,15 @@ class DSAIndexer(MegatronModule): + # Compute attention scores: q @ k^T + # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T + # -> [seqlen_q, batch, index_n_heads, seqlen_k] +- index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) ++ cp_size = parallel_state.get_context_parallel_world_size() ++ if cp_size == 1: ++ index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) ++ else: ++ # because k is small (only 1 head), do just one all_gather ++ k_buffer = torch.cat(torch.distributed.nn.functional.all_gather(k, group=self.pg_collection.cp), dim=0) # k_buffer: [[chunk_0, chunk_3, chunk_1, chunk_2], batch, index_head_dim] ++ index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k_buffer.float()) # [s_q_local, batch, index_n_heads, s_k_global] ++ # rank 0: q [chunk_0, chunk_3], k[chunk_0, chunk_3, chunk_1, chunk_2] ++ # rank 1: q [chunk_1, chunk_2], k[chunk_0, chunk_3, chunk_1, chunk_2] + + # Apply ReLU activation. + index_scores = torch.relu(index_scores) +@@ -606,7 +696,10 @@ class DSAIndexer(MegatronModule): + # ========================================= + # Select top-k indices + # ========================================= +- topk_k = min(self.index_topk, seqlen) ++ cp_size = parallel_state.get_context_parallel_world_size() ++ ++ seqlen_k_global = k.shape[0] * cp_size ++ topk_k = min(self.index_topk, seqlen_k_global) + # [batch, seqlen, index_topk] + topk_indices = index_scores.topk(topk_k, dim=-1)[1] + +@@ -687,6 +780,57 @@ def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): + output = output.reshape(sq, b, np * hnv) + return output + ++def get_causal_mask(sq, skv, device): ++ cp_size = parallel_state.get_context_parallel_world_size() ++ cp_rank = parallel_state.get_context_parallel_rank() ++ skv_global = skv * cp_size ++ ++ if cp_size == 1: ++ causal_mask = torch.triu( ++ torch.ones((sq, skv), dtype=torch.bool, device=device), ++ diagonal=1, ++ ) ++ else: ++ sq_half = sq // 2 ++ global_q_positions = torch.cat([ ++ torch.arange(cp_rank * sq_half, (cp_rank + 1) * sq_half, device=device), ++ torch.arange(skv_global - (cp_rank + 1) * sq_half, skv_global - cp_rank * sq_half, device=device) ++ ]) ++ ++ global_k_positions = torch.arange(skv_global, device=device) ++ # [sq, 1] < [1, skv_global] -> [sq, skv_global] ++ causal_mask = global_q_positions.unsqueeze(1) < global_k_positions.unsqueeze(0) ++ # convert to zz mask ++ chunked = causal_mask.chunk(dim=1, chunks=cp_size * 2) ++ causal_mask = [_x for _p in zip(chunked[:cp_size], reversed(chunked[cp_size:])) for _x in _p] ++ causal_mask = torch.cat(causal_mask, dim=1) ++ ++ return causal_mask ++ ++def unfused_dsa_fn_with_cp(query, key, value, topk_indices, softmax_scale): ++ pg = parallel_state.get_context_parallel_group() ++ cp_size = parallel_state.get_context_parallel_world_size() ++ cp_rank = parallel_state.get_context_parallel_rank() ++ ++ sq, b, np, hn = query.size() ++ skv = key.size(0) ++ hnv = value.size(3) ++ ++ skv_global = skv * cp_size ++ ++ sparse_mask = torch.ones((b, sq, skv_global), dtype=torch.bool, device=query.device) ++ sparse_mask.scatter_(-1, topk_indices, False) ++ ++ causal_mask = get_causal_mask(sq, skv, query.device) ++ ++ combined_mask = sparse_mask | causal_mask.unsqueeze(0) ++ ++ attention_mask_for_cp = combined_mask.unsqueeze(1) # [b, 1, sq, skv_global] ++ output = AttentionFuncionWithContextParallel.apply( ++ query, key, value, attention_mask_for_cp, 0.0, softmax_scale, pg, True ++ ) ++ return output.reshape(sq, b, np * hnv) ++ + + class DSAttention(MegatronModule): + """ +@@ -768,18 +912,17 @@ class DSAttention(MegatronModule): + # Generate upper triangular mask with -inf above diagonal, 0 elsewhere + # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal) + # float_mask [sq, skv] +- float_mask = torch.triu( +- torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device), +- diagonal=1, +- ) ++ mask = get_causal_mask(sq, skv, x.device) + else: +- assert attention_mask.shape == (b, 1, sq, skv), 'attention_mask shape mismatch' ++ skv_global = skv * parallel_state.get_context_parallel_world_size() ++ assert attention_mask.shape == (b, 1, sq, skv_global), 'attention_mask shape mismatch' + # [b, 1, sq, skv] -> [b, sq, skv] + mask = attention_mask.squeeze() +- # float_mask [b, sq, skv] +- float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( +- mask, float('-inf') +- ) ++ ++ # float_mask [b, sq, skv] ++ float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( ++ mask, float('-inf') ++ ) + + # =================================== + # Get index scores and top-k indices +@@ -791,7 +934,7 @@ class DSAttention(MegatronModule): + # =================================== + # Run sparse attention kernel + # =================================== +- output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) ++ output = unfused_dsa_fn_with_cp(query, key, value, topk_indices, self.softmax_scale) + + # =================================== + # Attach indexer loss +diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py +index 3953d933b..0ec5029dd 100644 +--- a/megatron/core/transformer/multi_latent_attention.py ++++ b/megatron/core/transformer/multi_latent_attention.py +@@ -6,6 +6,7 @@ from dataclasses import dataclass + from typing import NoReturn, Optional, Union + + import torch ++import torch.nn.functional as F + + try: + from einops import rearrange +@@ -198,6 +199,64 @@ class MultiLatentAttention(Attention): + # the quantized tensor. + set_save_original_input(self.linear_proj) + ++ def convert_thd_and_bsnh(self, src, packed_seq_params, to_bsd): ++ pg = parallel_state.get_context_parallel_group() ++ cp_size = parallel_state.get_context_parallel_world_size() ++ cp_rank = parallel_state.get_context_parallel_rank() ++ ++ seq_len_global = packed_seq_params.max_seqlen_q ++ seq_len_local = seq_len_global // cp_size ++ cu_seqlens_local = packed_seq_params.cu_seqlens_q // cp_size ++ b = len(packed_seq_params.cu_seqlens_q) - 1 ++ t = cu_seqlens_local[-1].item() ++ d = src.shape[-1] ++ ++ if to_bsd: ++ dst = torch.zeros(seq_len_local, b, d, ++ device=src.device, dtype=src.dtype) ++ else: ++ dst = torch.empty((t, 1, d), device=src.device, dtype=src.dtype) ++ ++ if cp_size == 1: ++ for i in range(b): ++ start, end = cu_seqlens_local[i].item(), cu_seqlens_local[i+1].item() ++ if to_bsd: ++ dst[:end-start, i] = src[start:end, 0] ++ else: ++ dst[start:end, 0] = src[:end-start, i] ++ else: ++ gathered = torch.stack( # TODO, may be too large? largest size: cp_size * s * b * h ++ torch.distributed.nn.functional.all_gather(src, group=pg), dim=0 ++ ) ++ for i in range(b): ++ start, end = cu_seqlens_local[i].item(), cu_seqlens_local[i+1].item() ++ len_i = end - start ++ half_len_i = len_i // 2 ++ half = start + half_len_i ++ chunk_size = seq_len_local // 2 ++ s1, e1 = chunk_size * cp_rank, chunk_size * (cp_rank + 1) ++ s2, e2 = chunk_size * (2 * cp_size - cp_rank - 1), chunk_size * (2 * cp_size - cp_rank) ++ ++ if to_bsd: ++ first_half = gathered[:, start:half, 0].contiguous().view(cp_size * half_len_i, -1) ++ second_half = gathered[:, half:end, 0].flip(dims=[0]).contiguous().view(cp_size * half_len_i, -1) ++ padded = F.pad( ++ torch.cat([first_half, second_half], dim=0), ++ (0, 0, 0, seq_len_global - cp_size * len_i), value=0 ++ ) ++ dst[:, i] = torch.cat([padded[s1:e1], padded[s2:e2]], dim=0) ++ else: ++ first_chunk = gathered[:, :chunk_size, i] # s1, s2, ... ++ second_chunk = gathered[:, chunk_size:seq_len_local, i].flip(dims=[0]) # s_n, s_n-1 ... ++ ++ full_padded = torch.cat([first_chunk, second_chunk], dim=0).contiguous().view(seq_len_global, d) ++ ++ ++ dst[start:half, 0] = full_padded[half_len_i * cp_rank:half_len_i * (cp_rank + 1)] ++ dst[half:end, 0] = full_padded[half_len_i * (2 * cp_size - cp_rank - 1):half_len_i * (2 * cp_size - cp_rank)] ++ ++ return dst ++ + def forward( + self, + hidden_states, +@@ -237,6 +296,13 @@ class MultiLatentAttention(Attention): + if self.config.cache_mla_latents: + self.prepare_for_absorption() + ++ original_packed_seq_params = None ++ if (self.config.experimental_attention_variant == "dsa" and ++ packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'): ++ original_packed_seq_params = packed_seq_params ++ hidden_states = self.convert_thd_and_bsnh(hidden_states, packed_seq_params, to_bsd=True) ++ packed_seq_params = None ++ + # ===================== + # Query, Key, and Value + # ===================== +@@ -306,8 +372,6 @@ class MultiLatentAttention(Attention): + attn_mask_type=attn_mask_type, + ) + elif self.config.experimental_attention_variant == "dsa": +- # For dsa we need to pass in the original hidden states and the compressed +- # query representation. + core_attn_out = self.core_attention( + query, + key, +@@ -358,11 +422,9 @@ class MultiLatentAttention(Attention): + # Flatten back: [seq, batch, num_heads * v_head_dim] + core_attn_out = core_attn_out.view(core_attn_out.size(0), core_attn_out.size(1), -1) + +- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': +- # reshape to same output shape as unpacked case +- # (t, np, hn) -> (t, b=1, h=np*hn) +- # t is the pack size = sum (sq_i) +- # note that batch is a dummy dimension in the packed case ++ if original_packed_seq_params is not None: ++ core_attn_out = self.convert_thd_and_bsnh(core_attn_out, original_packed_seq_params, to_bsd=False) ++ elif packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) + + if self.recompute_up_proj: +diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py +index a3a167549..98391fda6 100644 +--- a/megatron/core/transformer/transformer_config.py ++++ b/megatron/core/transformer/transformer_config.py +@@ -918,9 +918,9 @@ class TransformerConfig(ModelParallelConfig): + f" but got {self.context_parallel_size=}." + ) + elif self.experimental_attention_variant == "dsa": +- assert ( +- self.context_parallel_size == 1 +- ), "Currently context parallelism is not supported by DSAttention!" ++ # assert ( ++ # self.context_parallel_size == 1 ++ # ), "Currently context parallelism is not supported by DSAttention!" + assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" + + if self.fp8: diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index 41495f3969b..c7649cd8b83 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -42,7 +42,8 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio # TODO: find a clearer way. if converted_name.endswith("_scale"): continue - quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size)) + if_use_ue8m0_in_moe = True if args.sglang_moe_a2a_backend == "deepep" else False + quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size, if_use_ue8m0_in_moe=if_use_ue8m0_in_moe)) return quantize_named_params @@ -83,14 +84,14 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio return converted_named_params -def _quantize_param(name, weight, weight_block_size): +def _quantize_param(name, weight, weight_block_size, if_use_ue8m0_in_moe=True): assert name.endswith(".weight"), f"Expected weight parameter, got {name}" FP8_MIN = torch.finfo(torch.float8_e4m3fn).min FP8_MAX = torch.finfo(torch.float8_e4m3fn).max if weight_block_size is not None: if should_deepgemm_weight_requant_ue8m0 and should_deepgemm_weight_requant_ue8m0( weight_block_size=weight_block_size - ): + ) and if_use_ue8m0_in_moe: qweight, scale = quant_weight_ue8m0(weight, weight_block_size=weight_block_size) scale = transform_scale_ue8m0(scale, mn=qweight.shape[-2]) else: diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index a7f88d13762..abc790761d0 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -858,7 +858,7 @@ def loss_function( return ( loss, - torch.tensor(num_tokens if args.calculate_per_token_loss else 1, device=logits.device), + torch.tensor(num_tokens if args.calculate_per_token_loss else 1, dtype=torch.int, device=logits.device), { "keys": list(log.keys()), "values": torch.tensor( From 1a25680d7fb53ab26b076f5ed398c6e0ed1305e8 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sat, 13 Dec 2025 19:13:50 -0800 Subject: [PATCH 04/30] update dockerfile: TE version, fast-hadamard-transform --- docker/deepseekv32/Dockerfile | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile index 5787b1ce478..817bbbb31c8 100644 --- a/docker/deepseekv32/Dockerfile +++ b/docker/deepseekv32/Dockerfile @@ -38,12 +38,18 @@ RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f RUN pip install flash-linear-attention==0.4.0 +RUN git clone https://github.com/Dao-AILab/fast-hadamard-transform.git fast-hadamard-transform && \ + cd fast-hadamard-transform && \ + pip install -v . --no-build-isolation && \ + cd /root && \ + rm -rf fast-hadamard-transform + # TE does not have wheel on cuda 13 yet, thus need to install from source RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ pip install nvidia-mathdx==25.6.0 && \ pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.8; \ else \ - pip -v install --no-build-isolation "transformer_engine[pytorch]==2.8.0"; \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ fi RUN NVCC_APPEND_FLAGS="--threads 4" \ From f1674e9c06cc863b9fc10cd5599ada60431723bc Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sat, 13 Dec 2025 19:36:27 -0800 Subject: [PATCH 05/30] update patches --- docker/deepseekv32/Dockerfile | 16 +++++- .../{megatron_dsv32.patch => megatron.patch} | 0 docker/deepseekv32/transformers.patch | 21 ++++++++ miles/backends/megatron_utils/actor.py | 3 -- miles/utils/deepseek_v32_patch.py | 50 ------------------- miles_plugins/mbridge/__init__.py | 3 -- 6 files changed, 36 insertions(+), 57 deletions(-) rename docker/deepseekv32/{megatron_dsv32.patch => megatron.patch} (100%) create mode 100644 docker/deepseekv32/transformers.patch delete mode 100644 miles/utils/deepseek_v32_patch.py diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile index 817bbbb31c8..5692d0e256e 100644 --- a/docker/deepseekv32/Dockerfile +++ b/docker/deepseekv32/Dockerfile @@ -61,6 +61,10 @@ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \ pip install -e . +RUN git clone https://github.com/huggingface/transformers.git && \ + cd transformers && git checkout 40dc11cd3eb4126652aa41ef8272525affd4a636 && \ + pip install -e . + RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall RUN pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation @@ -102,7 +106,7 @@ RUN [ ! -f /root/.tmux.conf ] || rm /root/.tmux.conf # ====================================== Patches ============================================ -COPY docker/deepseekv32/${PATCH_VERSION}/megatron_dsv32.patch /root/Megatron-LM/ +COPY docker/deepseekv32/megatron.patch /root/Megatron-LM/ RUN cd Megatron-LM && \ git update-index --refresh && \ git apply megatron.patch --3way && \ @@ -112,6 +116,16 @@ RUN cd Megatron-LM && \ fi && \ rm megatron.patch +COPY docker/deepseekv32/transformers.patch /root/transformers/ +RUN cd transformers && \ + git update-index --refresh && \ + git apply transformers.patch --3way && \ + if grep -R -n '^<<<<<<< ' .; then \ + echo "Patch failed to apply cleanly. Please resolve conflicts." && \ + exit 1; \ + fi && \ + rm transformers.patch + # TODO temporarily skip patching for GB200/GB300 (and require users to bring their own sglang version). should add back later. COPY docker/patch/${PATCH_VERSION}/sglang.patch /sgl-workspace/sglang/ RUN if [ "$ENABLE_SGLANG_PATCH" = "1" ]; then \ diff --git a/docker/deepseekv32/megatron_dsv32.patch b/docker/deepseekv32/megatron.patch similarity index 100% rename from docker/deepseekv32/megatron_dsv32.patch rename to docker/deepseekv32/megatron.patch diff --git a/docker/deepseekv32/transformers.patch b/docker/deepseekv32/transformers.patch new file mode 100644 index 00000000000..61bc7b48306 --- /dev/null +++ b/docker/deepseekv32/transformers.patch @@ -0,0 +1,21 @@ +diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py +index 281bb0e773..6b8ae9f843 100644 +--- a/src/transformers/models/auto/configuration_auto.py ++++ b/src/transformers/models/auto/configuration_auto.py +@@ -1330,6 +1330,16 @@ class AutoConfig: + ) + config_dict["model_type"] = "ministral" + ++ if config_dict["model_type"] == "deepseek_v32": ++ logger.info( ++ "Detected deepseek_v32 model, treating as deepseek_v3 for compatibility." ++ ) ++ config_dict["model_type"] = "deepseek_v3" ++ if "architectures" in config_dict: ++ config_dict["architectures"] = [ ++ arch.replace("DeepseekV32", "DeepseekV3") for arch in config_dict["architectures"] ++ ] ++ + try: + config_class = CONFIG_MAPPING[config_dict["model_type"]] + except KeyError: diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index a12d1be4c95..a92198a6744 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -13,9 +13,6 @@ from torch_memory_saver import torch_memory_saver from transformers import AutoConfig, AutoTokenizer -from miles.utils.deepseek_v32_patch import apply_deepseek_v32_patch -apply_deepseek_v32_patch() - from miles.ray.train_actor import TrainRayActor from miles.utils import train_dump_utils from miles.utils.context_utils import with_defer diff --git a/miles/utils/deepseek_v32_patch.py b/miles/utils/deepseek_v32_patch.py deleted file mode 100644 index 94009070650..00000000000 --- a/miles/utils/deepseek_v32_patch.py +++ /dev/null @@ -1,50 +0,0 @@ -import os -import json -import tempfile -from transformers import AutoConfig - -_patched = False - - -def apply_deepseek_v32_patch(restore_model_type=False): - global _patched - - if _patched: - return - - _original_from_pretrained = AutoConfig.from_pretrained - - def _patched_from_pretrained(pretrained_model_name_or_path, *args, **kwargs): - if isinstance(pretrained_model_name_or_path, str) and os.path.isdir(pretrained_model_name_or_path): - config_file = os.path.join(pretrained_model_name_or_path, "config.json") - if os.path.exists(config_file): - try: - with open(config_file, "r") as f: - config_json = json.load(f) - - if config_json.get("model_type") == "deepseek_v32": - config_json["model_type"] = "deepseek_v3" - if "architectures" in config_json: - config_json["architectures"] = ["DeepseekV3ForCausalLM"] - - tmp_path = os.path.join(tempfile.gettempdir(), "_tmp_config_folder") - os.makedirs(tmp_path, exist_ok=True) - unique_path = os.path.join(tmp_path, f"deepseek_v32_{os.getpid()}.json") - - with open(unique_path, "w") as f: - json.dump(config_json, f) - - config = _original_from_pretrained(unique_path, *args, **kwargs) - - if restore_model_type: - object.__setattr__(config, "model_type", "deepseek_v32") - - return config - except Exception: - pass - - return _original_from_pretrained(pretrained_model_name_or_path, *args, **kwargs) - - AutoConfig.from_pretrained = _patched_from_pretrained - _patched = True - diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index 0e259d6b0f5..77741cb17bc 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -2,9 +2,6 @@ import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) -from miles.utils.deepseek_v32_patch import apply_deepseek_v32_patch -apply_deepseek_v32_patch(restore_model_type=True) - from .deepseekv32 import DeepseekV32Bridge from .glm4 import GLM4Bridge from .glm4moe import GLM4MoEBridge From 66d3b24299a0592cf7d6400ff762984deb80f8a7 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sat, 13 Dec 2025 19:46:26 -0800 Subject: [PATCH 06/30] update script --- scripts/run_deepseek_v3.2_5layer.py | 13 +-- scripts/train_dsv32.py | 121 ---------------------------- 2 files changed, 7 insertions(+), 127 deletions(-) delete mode 100644 scripts/train_dsv32.py diff --git a/scripts/run_deepseek_v3.2_5layer.py b/scripts/run_deepseek_v3.2_5layer.py index bcbb3a891af..a6a841a8d62 100644 --- a/scripts/run_deepseek_v3.2_5layer.py +++ b/scripts/run_deepseek_v3.2_5layer.py @@ -19,7 +19,7 @@ class ScriptArgs(U.ExecuteTrainConfig): enable_deepep: bool = False extra_args: str = "" task: Literal["dapo_aime", "gsm8k"] = "dapo_aime" - mode: Literal["normal", "debug_minimal"] = "debug_minimal" + mode: Literal["normal", "debug_minimal"] = "normal" @app.command() @@ -63,7 +63,7 @@ def train(args: ScriptArgs): rollout_args += ( "--prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl " "--input-key prompt " - f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 32768} " + f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} " ) eval_args += ( "--eval-prompt-data aime /root/aime-2024/aime-2024.jsonl " @@ -87,8 +87,8 @@ def train(args: ScriptArgs): "--tensor-model-parallel-size 1 " "--sequence-parallel " "--pipeline-model-parallel-size 1 " - "--context-parallel-size 1 " - "--expert-model-parallel-size 4 " + "--context-parallel-size 8 " + f"--expert-model-parallel-size {args.num_gpus_per_node} " "--expert-tensor-parallel-size 1 " ) elif args.num_nodes <= 4: @@ -136,7 +136,7 @@ def train(args: ScriptArgs): ) sglang_decode_max_bs = 256 - sglang_world_size = 4 if args.num_nodes <= 4 else 64 + sglang_world_size = args.num_gpus_per_node if args.num_nodes <= 4 else 64 sglang_attn_dp_size = 1 if args.num_nodes <= 4 else 8 sglang_attn_tp_size = sglang_world_size // sglang_attn_dp_size sglang_args = ( @@ -178,6 +178,7 @@ def train(args: ScriptArgs): f"--dump-details /root/shared_data/{args.run_id}/dump_details " "--disable-weights-backuper " "--model-name deepseekv32 " + "--train-memory-margin-bytes 1073741824 " ) train_args = ( @@ -195,7 +196,7 @@ def train(args: ScriptArgs): U.execute_train( train_args=train_args, - train_script="scripts/train_dsv32.py", + train_script="train.py", config=args, num_gpus_per_node=args.num_gpus_per_node, megatron_model_type="deepseek-v32-5layer", diff --git a/scripts/train_dsv32.py b/scripts/train_dsv32.py deleted file mode 100644 index 848cef803fb..00000000000 --- a/scripts/train_dsv32.py +++ /dev/null @@ -1,121 +0,0 @@ -import sys -import os -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) - -from miles.utils.deepseek_v32_patch import apply_deepseek_v32_patch -apply_deepseek_v32_patch() - -from turtle import mode -import ray -from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS -from typing import Optional - -try: - from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH -except ImportError: - GPU_MEMORY_TYPE_CUDA_GRAPH = None - -from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models -from miles.utils.arguments import parse_args -from miles.utils.logging_utils import configure_logger -from miles.utils.tracking_utils import init_tracking - - -def train(args): - configure_logger() - # allocate the GPUs - pgs = create_placement_groups(args) - init_tracking(args) - - # create the rollout manager, with sglang engines inside. - # need to initialize rollout manager first to calculate num_rollout - rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) - - # create the actor and critic models - actor_model, critic_model = create_training_models(args, pgs, rollout_manager) - - if args.offload_rollout: - ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])) - - # always update weight first so that sglang has the loaded weights from training. - actor_model.update_weights() - - if args.check_weight_update_equal: - ray.get(rollout_manager.check_weights.remote(action="compare")) - - if args.offload_rollout: - if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: - ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])) - ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE])) - - # special case for eval-only - if args.num_rollout == 0 and args.eval_interval is not None: - ray.get(rollout_manager.eval.remote(rollout_id=0)) - - def offload_train(): - if args.offload_train: - if args.use_critic: - critic_model.offload() - if rollout_id >= args.num_critic_only_steps: - actor_model.offload() - else: - actor_model.offload() - else: - actor_model.clear_memory() - - def onload_rollout(): - if args.offload_rollout: - ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])) - - # train loop. - # note that for async training, one can change the position of the sync operation(ray.get). - for rollout_id in range(args.start_rollout_id, args.num_rollout): - # TODO extract the duplicated eval logic - if args.eval_interval is not None and rollout_id == 0: - ray.get(rollout_manager.eval.remote(rollout_id)) - - rollout_data_ref = ray.get(rollout_manager.generate.remote(rollout_id)) - - if args.offload_rollout: - ray.get(rollout_manager.offload.remote()) - - if args.use_critic: - critic_train_handle = critic_model.async_train(rollout_id, rollout_data_ref) - if rollout_id >= args.num_critic_only_steps: - ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) - ray.get(critic_train_handle) - else: - ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) - - if args.save_interval is not None and ( - (rollout_id + 1) % args.save_interval == 0 - or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0) - ): - if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps): - actor_model.save_model(rollout_id) - if args.use_critic: - critic_model.save_model(rollout_id) - if args.rollout_global_dataset: - ray.get(rollout_manager.save.remote(rollout_id)) - - offload_train() - onload_rollout() - actor_model.update_weights() - - if args.offload_rollout: - if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: - ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])) - ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE])) - - if args.eval_interval is not None and ( - (rollout_id + 1) % args.eval_interval == 0 - or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0) - ): - ray.get(rollout_manager.eval.remote(rollout_id)) - - ray.get(rollout_manager.dispose.remote()) - - -if __name__ == "__main__": - args = parse_args() - train(args) From ccdff922d7c21c03a3a4d7d7824f2f3a4ac623d0 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 14 Dec 2025 16:26:39 -0800 Subject: [PATCH 07/30] minor fix --- docker/deepseekv32/megatron.patch | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index 5ad0d3afe08..7a21a3c65d8 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -1,5 +1,5 @@ diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py -index 89659a1d7..f1d6855ee 100644 +index 89659a1d7..38efa896c 100644 --- a/megatron/core/transformer/dot_product_attention_context_parallel.py +++ b/megatron/core/transformer/dot_product_attention_context_parallel.py @@ -132,10 +132,10 @@ class AllGatherComm: @@ -15,15 +15,6 @@ index 89659a1d7..f1d6855ee 100644 zz_mask = attention_mask else: chunked = attention_mask.chunk(dim=3, chunks=cp_size * 2) -@@ -143,7 +143,7 @@ def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stri - zz_mask = torch.cat(zz_mask, dim=3) - attn_bias = torch.zeros(zz_mask.shape, device=device, dtype=dtype) - attn_bias.masked_fill_(zz_mask, float('-inf')) -- attn_bias = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) -+ attn_bia = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) - return attn_bias - - @@ -151,7 +151,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): """Native attention function with context parallelism.""" From fd6bea68c9e89dc5e41f352eaaf6760d2b72a33d Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 22 Dec 2025 23:06:48 -0800 Subject: [PATCH 08/30] fix --- .../run-qwen3-4b-mis.sh | 20 +- .../processors/quantizer_fp8.py | 3 + .../megatron_utils/update_weight/common.py | 5 + miles/utils/external_utils/command_utils.py | 8 +- miles_plugins/mbridge/__init__.py | 2 +- miles_plugins/mbridge/deepseekv32.py | 7 + scripts/models/deepseek-v32-5layer.sh | 1 + scripts/models/deepseek-v32.sh | 69 ++++ scripts/run_deepseek_v3.2_5layer.py | 2 +- scripts/run_deepseek_v32.py | 295 ++++++++++++++++++ 10 files changed, 397 insertions(+), 15 deletions(-) create mode 100644 scripts/models/deepseek-v32-5layer.sh create mode 100644 scripts/models/deepseek-v32.sh create mode 100644 scripts/run_deepseek_v32.py diff --git a/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh b/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh index 300e8ac75b1..a130caa58f5 100644 --- a/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh +++ b/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh @@ -24,37 +24,37 @@ fi echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "/root/miles/scripts/models/qwen3-4B.sh" +source "/host_home/primary_synced/miles/scripts/models/qwen3-4B.sh" CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B + --hf-checkpoint /host_home/models/Qwen3-4B #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/Qwen3-4B_torch_dist + --ref-load /root/models/Qwen3-4B_torch_dist # --load /root/Qwen3-4B_miles/ - --save /root/Qwen3-4B_miles/ + --save /root/models/Qwen3-4B_miles/ --save-interval 200 ) ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --prompt-data /host_home/data/dapo-math-17k/dapo-math-17k.jsonl --input-key prompt --label-key label --apply-chat-template --rollout-shuffle --rm-type deepscaler --num-rollout 3000 - --rollout-batch-size 32 + --rollout-batch-size 8 --n-samples-per-prompt 8 --rollout-max-response-len 8192 --rollout-temperature 1 - --global-batch-size 256 + --global-batch-size 64 --balance-data ) EVAL_ARGS=( # --eval-interval 20 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --eval-prompt-data aime /host_home/data/aime-2024/aime-2024.jsonl --n-samples-per-eval-prompt 1 --eval-max-response-len 16384 --eval-top-p 1 @@ -127,7 +127,7 @@ CUSTOM_ARGS=( # launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 # Build the runtime environment JSON with proper variable substitution RUNTIME_ENV_JSON="{ @@ -142,7 +142,7 @@ ray job submit --address="http://127.0.0.1:8265" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ -- python3 train.py \ --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ + --actor-num-gpus-per-node 4 \ --colocate \ ${MODEL_ARGS[@]} \ ${CKPT_ARGS[@]} \ diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index c7649cd8b83..54bb1e67646 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -73,6 +73,9 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio "self_attention.linear_q_up_proj.weight", "self_attention.linear_kv_down_proj.weight", "self_attention.linear_kv_up_proj.weight", + # dsa indexer + "self_attention.core_attention.indexer.linear_wq_b.weight", + "self_attention.core_attention.indexer.linear_wk.weight", ]: quantize_named_params = [] for converted_name, param in converted_named_params: diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index 85fe76a1b8d..558a2e06f0c 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -202,6 +202,11 @@ def _named_params_and_buffers_global( expert_idx = int(expert_idx) + expert_offset yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param continue + + # TODO: a hacking here, need to be cleaner + duplicated = ['indexer.linear_weights_proj', 'indexer.linear_wk', 'indexer.linear_wq_b', 'linear_q_down_proj', 'linear_kv_down_proj'] + if any(dup in name for dup in duplicated): + param.parallel_mode = 'duplicated' layer_idx, rest = match.groups() layer_idx = int(layer_idx) + layer_offset diff --git a/miles/utils/external_utils/command_utils.py b/miles/utils/external_utils/command_utils.py index 8c7c9316b95..e6bf01a1e3f 100644 --- a/miles/utils/external_utils/command_utils.py +++ b/miles/utils/external_utils/command_utils.py @@ -51,14 +51,15 @@ def convert_checkpoint( exec_command( f"source {repo_base_dir}/scripts/models/{megatron_model_type}.sh && " - f"PYTHONPATH=/root/Megatron-LM " + # Use installed Megatron instead of hardcoded path + f"PYTHONPATH=/host_home/primary_synced/Megatron-LM " f"torchrun " f"--nproc-per-node {num_gpus_per_node} " f"{multinode_args}" f"tools/convert_hf_to_torch_dist.py " "${MODEL_ARGS[@]} " f"--hf-checkpoint {hf_checkpoint} " - f"--save {path_dst}" + f"--save {path_dst} " f"{extra_args}" ) @@ -139,7 +140,8 @@ def execute_train( runtime_env_json = json.dumps( { "env_vars": { - "PYTHONPATH": "/root/Megatron-LM/", + # Use installed Megatron instead of hardcoded path + "PYTHONPATH": "/host_home/primary_synced/Megatron-LM/", # If setting this in FSDP, the computation communication overlapping may have issues **( {} diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index 77741cb17bc..cc42522eeec 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -16,7 +16,7 @@ @classmethod def _patched_from_config(cls, hf_config, **kwargs): - if hf_config.model_type == "deepseek_v32": + if hasattr(hf_config, 'index_n_heads'): from mbridge.core.bridge import _MODEL_REGISTRY return _MODEL_REGISTRY['deepseek_v32'](hf_config, **kwargs) diff --git a/miles_plugins/mbridge/deepseekv32.py b/miles_plugins/mbridge/deepseekv32.py index fb9355f5ee3..aae07ee5321 100644 --- a/miles_plugins/mbridge/deepseekv32.py +++ b/miles_plugins/mbridge/deepseekv32.py @@ -6,6 +6,13 @@ @register_model("deepseek_v32") class DeepseekV32Bridge(DeepseekV3Bridge): + # Weights with parallel_mode="duplicated" that should NOT be gathered across TP + _DUPLICATED_WEIGHTS = { + "self_attention.core_attention.indexer.linear_wq_b.weight", + "self_attention.core_attention.indexer.linear_wk.weight", + "self_attention.core_attention.indexer.linear_weights_proj.weight", + } + _ATTENTION_MAPPING = ( DeepseekV3Bridge._ATTENTION_MAPPING.copy() ) diff --git a/scripts/models/deepseek-v32-5layer.sh b/scripts/models/deepseek-v32-5layer.sh new file mode 100644 index 00000000000..2466640afd5 --- /dev/null +++ b/scripts/models/deepseek-v32-5layer.sh @@ -0,0 +1 @@ +MODEL_ARGS_NUM_LAYERS=5 source "$(dirname -- "${BASH_SOURCE[0]}")/deepseek-v32.sh" diff --git a/scripts/models/deepseek-v32.sh b/scripts/models/deepseek-v32.sh new file mode 100644 index 00000000000..a98f2f561d1 --- /dev/null +++ b/scripts/models/deepseek-v32.sh @@ -0,0 +1,69 @@ +NLAYERS="${MODEL_ARGS_NUM_LAYERS:-61}" +FIRST_K_DENSE_REPLACE=3 + +arr=() +for ((i=0; i Date: Tue, 16 Dec 2025 23:11:30 +0000 Subject: [PATCH 09/30] update --- docker/deepseekv32/Dockerfile | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile index 5692d0e256e..a52b35a6738 100644 --- a/docker/deepseekv32/Dockerfile +++ b/docker/deepseekv32/Dockerfile @@ -6,7 +6,7 @@ FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang ARG PATCH_VERSION=latest ARG MEGATRON_COMMIT=436065a86b749ca3b50eebca68f55c9e690a9f63 -ARG ENABLE_CUDA_13=0 +ARG ENABLE_CUDA_13=1 ARG ENABLE_SGLANG_PATCH=0 @@ -47,7 +47,8 @@ RUN git clone https://github.com/Dao-AILab/fast-hadamard-transform.git fast-hada # TE does not have wheel on cuda 13 yet, thus need to install from source RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ pip install nvidia-mathdx==25.6.0 && \ - pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.8; \ + pip install pybind11 && \ + pip -v install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10; \ else \ pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ fi @@ -83,11 +84,6 @@ RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ python3 -m pip install https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu130-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps; \ fi -# This patch is merged into main, but we are using stable version, thus still need it -RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ - curl -L https://github.com/NVIDIA/TransformerEngine/pull/2286.patch -o /root/te2286.patch && (cd /usr/local/lib/python3.12/dist-packages/transformer_engine && (patch -p2 < /root/te2286.patch)); \ - fi - # AMEM # we need to create a fake libcuda.so.1 to make the linker happy when building AMEM ENV CUDA_DIR=/usr/local/cuda From bb09c276aabd7ad6959c792de5998f6efac52b33 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Fri, 19 Dec 2025 11:36:52 -0800 Subject: [PATCH 10/30] init --- miles/utils/arguments.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 79b2c419ca6..7aed05451bf 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1701,6 +1701,9 @@ def miles_validate_args(args): args.use_dynamic_batch_size is False ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." + if args.disable_thd_format: + assert args.train_backend == "megatron", "disable_thd_format is only supported for megatron backend." + def hf_validate_args(args, hf_config): def equal(x, y): From 8af1384dad48ad07628533ebf9b2b3682059c0cf Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 22 Dec 2025 20:28:44 -0800 Subject: [PATCH 11/30] supported bshd --- miles/backends/training_utils/data.py | 7 +++++++ miles/backends/training_utils/loss.py | 5 ++++- miles/utils/arguments.py | 5 +++-- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index 67bb30108d1..b8ce38c970b 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -134,6 +134,13 @@ def get_batch( tokens = [slice_with_cp(t, pad_token_id, parallel_state, qkv_format, max_seqlen) for t in tokens] tokens = torch.stack(tokens) + if qkv_format == "bshd": + max_seqlen = batch["max_seq_len"][0] + assert max([t.size(0) for t in tokens]) <= max_seqlen + tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] + tokens = torch.stack(tokens) + # TODO: padding to multiples? + elif qkv_format == "thd": tokens = [slice_with_cp(t, pad_token_id, parallel_state, qkv_format) for t in tokens] diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index abc790761d0..1752044fe30 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -87,6 +87,7 @@ def get_responses( tokens_chunk = tokens[-response_length:] else: # TODO: this is super ugly... do better abstraction. + _max_seq_len = max_seq_len[i] if max_seq_len is not None else None chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( total_length, response_length, parallel_state, qkv_format, max_seq_len ) @@ -101,7 +102,9 @@ def get_responses( tokens_1 = tokens[tokens_offset[1][0] : tokens_offset[1][1]] assert logits_0.size(0) == tokens_0.size(0), f"{logits_0.size(0)} vs {tokens_0.size(0)}" - assert logits_1.size(0) == tokens_1.size(0), f"{logits_1.size(0)} vs {tokens_1.size(0)}" + assert logits_1.size(0) == tokens_1.size(0), f"{logits_1.size(0)} vs {tokens_1.size(0)}, chunks_offset {chunks_offset}, logits_offset {logits_offset}, \ + tokens_offset {tokens_offset}, logits_1 range {(end + chunk_size, end + 2 * chunk_size)}, chunk_size {chunk_size}, max_seq_len {_max_seq_len}, \ + logits.shape {logits.shape}" logits_chunk = torch.cat([logits_0, logits_1], dim=0) tokens_chunk = torch.cat([tokens_0, tokens_1], dim=0) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7aed05451bf..7139f022087 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1701,8 +1701,9 @@ def miles_validate_args(args): args.use_dynamic_batch_size is False ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." - if args.disable_thd_format: - assert args.train_backend == "megatron", "disable_thd_format is only supported for megatron backend." + assert args.qkv_format in ['thd', 'bshd'], f"qkv_format {args.qkv_format} is not supported. (only 'thd' and 'bshd' are supported)" + if args.qkv_format == 'bshd': + assert args.train_backend == "megatron", "bshd format is only supported for megatron backend." def hf_validate_args(args, hf_config): From 42c680f224074e3f591b837e10a18eba46ade146 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 22 Dec 2025 21:08:10 -0800 Subject: [PATCH 12/30] lint --- miles/backends/training_utils/data.py | 4 ++-- miles/backends/training_utils/loss.py | 4 +--- miles/utils/arguments.py | 7 +++++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index b8ce38c970b..4ce8ec4682a 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -138,9 +138,9 @@ def get_batch( max_seqlen = batch["max_seq_len"][0] assert max([t.size(0) for t in tokens]) <= max_seqlen tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] - tokens = torch.stack(tokens) + tokens = torch.stack(tokens) # TODO: padding to multiples? - + elif qkv_format == "thd": tokens = [slice_with_cp(t, pad_token_id, parallel_state, qkv_format) for t in tokens] diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index 1752044fe30..72f7526ed0e 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -102,9 +102,7 @@ def get_responses( tokens_1 = tokens[tokens_offset[1][0] : tokens_offset[1][1]] assert logits_0.size(0) == tokens_0.size(0), f"{logits_0.size(0)} vs {tokens_0.size(0)}" - assert logits_1.size(0) == tokens_1.size(0), f"{logits_1.size(0)} vs {tokens_1.size(0)}, chunks_offset {chunks_offset}, logits_offset {logits_offset}, \ - tokens_offset {tokens_offset}, logits_1 range {(end + chunk_size, end + 2 * chunk_size)}, chunk_size {chunk_size}, max_seq_len {_max_seq_len}, \ - logits.shape {logits.shape}" + assert logits_1.size(0) == tokens_1.size(0), f"{logits_1.size(0)} vs {tokens_1.size(0)}" logits_chunk = torch.cat([logits_0, logits_1], dim=0) tokens_chunk = torch.cat([tokens_0, tokens_1], dim=0) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7139f022087..de5378d55e7 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1701,8 +1701,11 @@ def miles_validate_args(args): args.use_dynamic_batch_size is False ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." - assert args.qkv_format in ['thd', 'bshd'], f"qkv_format {args.qkv_format} is not supported. (only 'thd' and 'bshd' are supported)" - if args.qkv_format == 'bshd': + assert args.qkv_format in [ + "thd", + "bshd", + ], f"qkv_format {args.qkv_format} is not supported. (only 'thd' and 'bshd' are supported)" + if args.qkv_format == "bshd": assert args.train_backend == "megatron", "bshd format is only supported for megatron backend." From 0791a9c8ae8d9eff7dd1037180ec365b93f4020f Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 22 Dec 2025 21:53:24 -0800 Subject: [PATCH 13/30] rename, add argument assert, lint --- miles/backends/training_utils/data.py | 2 +- miles/backends/training_utils/loss.py | 2 +- miles/utils/arguments.py | 3 +++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index 4ce8ec4682a..d38161e6070 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -135,7 +135,7 @@ def get_batch( tokens = torch.stack(tokens) if qkv_format == "bshd": - max_seqlen = batch["max_seq_len"][0] + max_seqlen = batch["max_seq_lens"][0] assert max([t.size(0) for t in tokens]) <= max_seqlen tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] tokens = torch.stack(tokens) diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index 72f7526ed0e..37e81eb533a 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -87,7 +87,7 @@ def get_responses( tokens_chunk = tokens[-response_length:] else: # TODO: this is super ugly... do better abstraction. - _max_seq_len = max_seq_len[i] if max_seq_len is not None else None + max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( total_length, response_length, parallel_state, qkv_format, max_seq_len ) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index de5378d55e7..bf7b0467e38 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1707,6 +1707,9 @@ def miles_validate_args(args): ], f"qkv_format {args.qkv_format} is not supported. (only 'thd' and 'bshd' are supported)" if args.qkv_format == "bshd": assert args.train_backend == "megatron", "bshd format is only supported for megatron backend." + assert ( + args.use_dynamic_batch_size is False + ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead." def hf_validate_args(args, hf_config): From d8cb73a2fb6c8c39601a748b9a423f3597db877e Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sat, 27 Dec 2025 23:31:30 -0800 Subject: [PATCH 14/30] tmp fix --- miles/backends/megatron_utils/actor.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index a92198a6744..95803c73f15 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -462,6 +462,9 @@ def update_weights(self) -> None: if self.args.offload_train: reload_process_groups() + if isinstance(num_new_engines, tuple): + num_new_engines = num_new_engines[0] + if num_new_engines > 0: self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock) dist.barrier(group=get_gloo_group()) From 7009e11a966d347907d8e132d8a7b4db00370758 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 01:15:04 -0800 Subject: [PATCH 15/30] update megatron patch --- docker/deepseekv32/megatron.patch | 105 +----------------------------- 1 file changed, 1 insertion(+), 104 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index 7a21a3c65d8..e6a8a34a561 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -410,7 +410,7 @@ index fc994490b..7bc9a485e 100644 # =================================== # Attach indexer loss diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py -index 3953d933b..0ec5029dd 100644 +index 3953d933b..84301ed54 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -6,6 +6,7 @@ from dataclasses import dataclass @@ -421,109 +421,6 @@ index 3953d933b..0ec5029dd 100644 try: from einops import rearrange -@@ -198,6 +199,64 @@ class MultiLatentAttention(Attention): - # the quantized tensor. - set_save_original_input(self.linear_proj) - -+ def convert_thd_and_bsnh(self, src, packed_seq_params, to_bsd): -+ pg = parallel_state.get_context_parallel_group() -+ cp_size = parallel_state.get_context_parallel_world_size() -+ cp_rank = parallel_state.get_context_parallel_rank() -+ -+ seq_len_global = packed_seq_params.max_seqlen_q -+ seq_len_local = seq_len_global // cp_size -+ cu_seqlens_local = packed_seq_params.cu_seqlens_q // cp_size -+ b = len(packed_seq_params.cu_seqlens_q) - 1 -+ t = cu_seqlens_local[-1].item() -+ d = src.shape[-1] -+ -+ if to_bsd: -+ dst = torch.zeros(seq_len_local, b, d, -+ device=src.device, dtype=src.dtype) -+ else: -+ dst = torch.empty((t, 1, d), device=src.device, dtype=src.dtype) -+ -+ if cp_size == 1: -+ for i in range(b): -+ start, end = cu_seqlens_local[i].item(), cu_seqlens_local[i+1].item() -+ if to_bsd: -+ dst[:end-start, i] = src[start:end, 0] -+ else: -+ dst[start:end, 0] = src[:end-start, i] -+ else: -+ gathered = torch.stack( # TODO, may be too large? largest size: cp_size * s * b * h -+ torch.distributed.nn.functional.all_gather(src, group=pg), dim=0 -+ ) -+ for i in range(b): -+ start, end = cu_seqlens_local[i].item(), cu_seqlens_local[i+1].item() -+ len_i = end - start -+ half_len_i = len_i // 2 -+ half = start + half_len_i -+ chunk_size = seq_len_local // 2 -+ s1, e1 = chunk_size * cp_rank, chunk_size * (cp_rank + 1) -+ s2, e2 = chunk_size * (2 * cp_size - cp_rank - 1), chunk_size * (2 * cp_size - cp_rank) -+ -+ if to_bsd: -+ first_half = gathered[:, start:half, 0].contiguous().view(cp_size * half_len_i, -1) -+ second_half = gathered[:, half:end, 0].flip(dims=[0]).contiguous().view(cp_size * half_len_i, -1) -+ padded = F.pad( -+ torch.cat([first_half, second_half], dim=0), -+ (0, 0, 0, seq_len_global - cp_size * len_i), value=0 -+ ) -+ dst[:, i] = torch.cat([padded[s1:e1], padded[s2:e2]], dim=0) -+ else: -+ first_chunk = gathered[:, :chunk_size, i] # s1, s2, ... -+ second_chunk = gathered[:, chunk_size:seq_len_local, i].flip(dims=[0]) # s_n, s_n-1 ... -+ -+ full_padded = torch.cat([first_chunk, second_chunk], dim=0).contiguous().view(seq_len_global, d) -+ -+ -+ dst[start:half, 0] = full_padded[half_len_i * cp_rank:half_len_i * (cp_rank + 1)] -+ dst[half:end, 0] = full_padded[half_len_i * (2 * cp_size - cp_rank - 1):half_len_i * (2 * cp_size - cp_rank)] -+ -+ return dst -+ - def forward( - self, - hidden_states, -@@ -237,6 +296,13 @@ class MultiLatentAttention(Attention): - if self.config.cache_mla_latents: - self.prepare_for_absorption() - -+ original_packed_seq_params = None -+ if (self.config.experimental_attention_variant == "dsa" and -+ packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'): -+ original_packed_seq_params = packed_seq_params -+ hidden_states = self.convert_thd_and_bsnh(hidden_states, packed_seq_params, to_bsd=True) -+ packed_seq_params = None -+ - # ===================== - # Query, Key, and Value - # ===================== -@@ -306,8 +372,6 @@ class MultiLatentAttention(Attention): - attn_mask_type=attn_mask_type, - ) - elif self.config.experimental_attention_variant == "dsa": -- # For dsa we need to pass in the original hidden states and the compressed -- # query representation. - core_attn_out = self.core_attention( - query, - key, -@@ -358,11 +422,9 @@ class MultiLatentAttention(Attention): - # Flatten back: [seq, batch, num_heads * v_head_dim] - core_attn_out = core_attn_out.view(core_attn_out.size(0), core_attn_out.size(1), -1) - -- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': -- # reshape to same output shape as unpacked case -- # (t, np, hn) -> (t, b=1, h=np*hn) -- # t is the pack size = sum (sq_i) -- # note that batch is a dummy dimension in the packed case -+ if original_packed_seq_params is not None: -+ core_attn_out = self.convert_thd_and_bsnh(core_attn_out, original_packed_seq_params, to_bsd=False) -+ elif packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) - - if self.recompute_up_proj: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index a3a167549..98391fda6 100644 --- a/megatron/core/transformer/transformer_config.py From 4499325a72349cdecf7179ac9228e07c27565ea4 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 01:16:23 -0800 Subject: [PATCH 16/30] update transformers patch --- docker/deepseekv32/Dockerfile | 2 +- docker/deepseekv32/transformers.patch | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile index a52b35a6738..e94e725edea 100644 --- a/docker/deepseekv32/Dockerfile +++ b/docker/deepseekv32/Dockerfile @@ -63,7 +63,7 @@ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ pip install -e . RUN git clone https://github.com/huggingface/transformers.git && \ - cd transformers && git checkout 40dc11cd3eb4126652aa41ef8272525affd4a636 && \ + cd transformers && git checkout 8cb5963cc22174954e7dca2c0a3320b7dc2f4edc && \ pip install -e . RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall diff --git a/docker/deepseekv32/transformers.patch b/docker/deepseekv32/transformers.patch index 61bc7b48306..a7631aa00c2 100644 --- a/docker/deepseekv32/transformers.patch +++ b/docker/deepseekv32/transformers.patch @@ -1,11 +1,11 @@ diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py -index 281bb0e773..6b8ae9f843 100644 +index f6a12e7cef..22129a86ee 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py -@@ -1330,6 +1330,16 @@ class AutoConfig: +@@ -1355,6 +1355,15 @@ class AutoConfig: + "Detected mistral model with layer_types, treating as ministral for alternating attention compatibility. " ) config_dict["model_type"] = "ministral" - + if config_dict["model_type"] == "deepseek_v32": + logger.info( + "Detected deepseek_v32 model, treating as deepseek_v3 for compatibility." @@ -15,7 +15,6 @@ index 281bb0e773..6b8ae9f843 100644 + config_dict["architectures"] = [ + arch.replace("DeepseekV32", "DeepseekV3") for arch in config_dict["architectures"] + ] -+ + try: config_class = CONFIG_MAPPING[config_dict["model_type"]] - except KeyError: From 6f1e1300372ef0cb3e655119ef91dc6b8e10d9d0 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 01:17:30 -0800 Subject: [PATCH 17/30] disable amem --- docker/deepseekv32/Dockerfile | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile index e94e725edea..5ee2cd5a497 100644 --- a/docker/deepseekv32/Dockerfile +++ b/docker/deepseekv32/Dockerfile @@ -86,17 +86,17 @@ RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ # AMEM # we need to create a fake libcuda.so.1 to make the linker happy when building AMEM -ENV CUDA_DIR=/usr/local/cuda -ENV CUDA_STUBS=${CUDA_DIR}/lib64/stubs -RUN ln -s ${CUDA_STUBS}/libcuda.so ${CUDA_STUBS}/libcuda.so.1 && \ - echo "${CUDA_STUBS}" > /etc/ld.so.conf.d/z-cuda-stubs.conf && \ - ldconfig -RUN git clone https://github.com/inclusionAI/asystem-amem.git && \ - cd asystem-amem && git checkout 6483bb17c9a98b51c3a94b7048467d5b50fbad4b && \ - git submodule init && git submodule update && \ - MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi/ ./build.sh && \ - mv /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2 /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2.bak && \ - cp -r third_party/nccl/build/lib/* /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/ +# ENV CUDA_DIR=/usr/local/cuda +# ENV CUDA_STUBS=${CUDA_DIR}/lib64/stubs +# RUN ln -s ${CUDA_STUBS}/libcuda.so ${CUDA_STUBS}/libcuda.so.1 && \ +# echo "${CUDA_STUBS}" > /etc/ld.so.conf.d/z-cuda-stubs.conf && \ +# ldconfig +# RUN git clone https://github.com/inclusionAI/asystem-amem.git && \ +# cd asystem-amem && git checkout 6483bb17c9a98b51c3a94b7048467d5b50fbad4b && \ +# git submodule init && git submodule update && \ +# MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi/ ./build.sh && \ +# mv /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2 /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2.bak && \ +# cp -r third_party/nccl/build/lib/* /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/ RUN [ ! -f /root/.tmux.conf ] || rm /root/.tmux.conf From cbd2e9f58e725b9e34a54eadc7a9c8b59b73f98b Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 01:18:15 -0800 Subject: [PATCH 18/30] add script --- scripts/run_deepseek_v32.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/scripts/run_deepseek_v32.py b/scripts/run_deepseek_v32.py index 5787f4ce2d9..86af2c42a9c 100644 --- a/scripts/run_deepseek_v32.py +++ b/scripts/run_deepseek_v32.py @@ -18,8 +18,8 @@ class ScriptArgs(U.ExecuteTrainConfig): mode: Literal["normal", "debug_minimal"] = "debug_minimal" run_id: str = U.create_run_id() model_org: str = "deepseek-ai" - model_name: Literal["DeepSeek-V3.2", "DeepSeek-V3.2-5layer"] = "DeepSeek-V3.2-5layer" - megatron_model_type: Literal["deepseek-v32", "deepseek-v32-5layer"] = "deepseek-v32-5layer" + model_name: Literal["DeepSeek-V3.2", "DeepSeek-V3.2-5layer"] = "DeepSeek-V3.2" + megatron_model_type: Literal["deepseek-v32", "deepseek-v32-5layer"] = "deepseek-v32" num_gpus_per_node: int = 4 enable_eval: bool = True extra_args: str = "" @@ -164,10 +164,10 @@ def train(args: ScriptArgs): else: # TODO choose a good config (currently randomly change to suit 64gpu) perf_args = ( - "--tensor-model-parallel-size 1 " + "--tensor-model-parallel-size 8 " "--sequence-parallel " f"--pipeline-model-parallel-size {1 if args.model_name == 'DeepSeek-V3.2-5layer' else 4} " - "--context-parallel-size 8 " + "--context-parallel-size 2 " "--expert-model-parallel-size 16 " "--expert-tensor-parallel-size 1 " ) @@ -179,7 +179,8 @@ def train(args: ScriptArgs): "--recompute-method uniform " "--recompute-num-layers 1 " # ------------ - "--use-dynamic-batch-size " + # "--use-dynamic-batch-size " + "--micro-batch-size 1 " # TODO temp use tiny value "--max-tokens-per-gpu 2048 " # "--max-tokens-per-gpu 16384 " @@ -266,6 +267,7 @@ def train(args: ScriptArgs): "--model-name deepseekv32 " # for mbridge load "--train-memory-margin-bytes 1073741824 " # "--check-weight-update-equal " + "--qkv-format bshd " ) train_args = ( From 9dc5258244a3f3c8782780c8457fb893e32b1704 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 19:44:37 -0800 Subject: [PATCH 19/30] update --- docker/deepseekv32/Dockerfile | 4 +- miles/utils/external_utils/command_utils.py | 10 ++-- scripts/run_deepseek_v32.py | 61 ++++++++++++++------- 3 files changed, 48 insertions(+), 27 deletions(-) diff --git a/docker/deepseekv32/Dockerfile b/docker/deepseekv32/Dockerfile index 5ee2cd5a497..38e9eadfafd 100644 --- a/docker/deepseekv32/Dockerfile +++ b/docker/deepseekv32/Dockerfile @@ -1,4 +1,4 @@ -ARG SGLANG_IMAGE_TAG=dev +ARG SGLANG_IMAGE_TAG=v0.5.6.post2 FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang # ======================================== Arguments ============================================= @@ -6,7 +6,7 @@ FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang ARG PATCH_VERSION=latest ARG MEGATRON_COMMIT=436065a86b749ca3b50eebca68f55c9e690a9f63 -ARG ENABLE_CUDA_13=1 +ARG ENABLE_CUDA_13=0 ARG ENABLE_SGLANG_PATCH=0 diff --git a/miles/utils/external_utils/command_utils.py b/miles/utils/external_utils/command_utils.py index e6bf01a1e3f..bdfc864b425 100644 --- a/miles/utils/external_utils/command_utils.py +++ b/miles/utils/external_utils/command_utils.py @@ -26,6 +26,7 @@ def convert_checkpoint( extra_args: str = "", dir_dst: str = "/root", hf_checkpoint: str | None = None, + megatron_path: str = "/host_home/primary_synced/Megatron-LM", ): hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}" @@ -52,7 +53,7 @@ def convert_checkpoint( exec_command( f"source {repo_base_dir}/scripts/models/{megatron_model_type}.sh && " # Use installed Megatron instead of hardcoded path - f"PYTHONPATH=/host_home/primary_synced/Megatron-LM " + f"PYTHONPATH={megatron_path} " f"torchrun " f"--nproc-per-node {num_gpus_per_node} " f"{multinode_args}" @@ -68,9 +69,9 @@ def rsync_simple(path_src: str, path_dst: str): exec_command(f"mkdir -p {path_dst} && rsync -a --info=progress2 {path_src}/ {path_dst}") -def hf_download_dataset(full_name: str): +def hf_download_dataset(full_name: str, data_dir: str = "/root/datasets"): _, partial_name = full_name.split("/") - exec_command(f"hf download --repo-type dataset {full_name} --local-dir /root/datasets/{partial_name}") + exec_command(f"hf download --repo-type dataset {full_name} --local-dir {data_dir}/{partial_name}") def fp8_cast_bf16(path_src, path_dst): @@ -99,6 +100,7 @@ def execute_train( before_ray_job_submit=None, extra_env_vars=None, config: ExecuteTrainConfig | None = None, + megatron_path: str = "/host_home/primary_synced/Megatron-LM", ): if extra_env_vars is None: extra_env_vars = {} @@ -141,7 +143,7 @@ def execute_train( { "env_vars": { # Use installed Megatron instead of hardcoded path - "PYTHONPATH": "/host_home/primary_synced/Megatron-LM/", + "PYTHONPATH": f"{megatron_path}", # If setting this in FSDP, the computation communication overlapping may have issues **( {} diff --git a/scripts/run_deepseek_v32.py b/scripts/run_deepseek_v32.py index 86af2c42a9c..0cbf9ca05d5 100644 --- a/scripts/run_deepseek_v32.py +++ b/scripts/run_deepseek_v32.py @@ -5,7 +5,7 @@ import re from dataclasses import dataclass from typing import Literal - +from pathlib import Path import typer import miles.utils.external_utils.command_utils as U @@ -25,6 +25,28 @@ class ScriptArgs(U.ExecuteTrainConfig): extra_args: str = "" task: Literal["dapo_aime", "gsm8k"] = "dapo_aime" enable_deepep: bool = True + data_dir: str = "/root" + model_dir: str = "/root/.cache/dsv32" + model_local_dir: str = "/root/.cache/dsv32" + save_dir: str = "/root/.cache/dsv32" + megatron_path: str = "/root/Megatron-LM" + + +@app.command() +@U.dataclass_cli +def prepare_single(args: ScriptArgs): + """This script only needs to be executed on one node.""" + match args.task: + case "dapo_aime": + U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + U.hf_download_dataset("zhuzilin/aime-2024", data_dir=args.data_dir) + case "gsm8k": + U.hf_download_dataset("zhuzilin/gsm8k", data_dir=args.data_dir) + + U.fp8_cast_bf16( + path_src=f"{args.model_dir}/{args.model_name}", + path_dst=f"{args.model_dir}/{args.model_name}-bf16/", + ) @app.command() @@ -34,11 +56,6 @@ def prepare_spmd(args: ScriptArgs): extra_args = "--tensor-model-parallel-size 1 " "--expert-tensor-parallel-size 1 " if args.num_nodes == 1 and args.model_name == "DeepSeek-V3.2-5layer": extra_args += "--pipeline-model-parallel-size 1 " "--expert-model-parallel-size 1 " - elif args.model_name == "DeepSeek-V3.2-20layer": - extra_args += ( - "--expert-model-parallel-size 4 " - # PP info will be auto determined by converter script - ) else: extra_args += ( "--pipeline-model-parallel-size 8 " @@ -49,12 +66,13 @@ def prepare_spmd(args: ScriptArgs): U.convert_checkpoint( model_name=args.model_name, - hf_checkpoint=f"/root/models/{args.model_name}-bf16", + hf_checkpoint=f"{args.model_dir}/{args.model_name}-bf16", megatron_model_type=args.megatron_model_type, num_gpus_per_node=args.num_gpus_per_node, - multinode=True, + multinode=True if args.num_nodes > 1 else False, extra_args=extra_args, - dir_dst="/root/models", + dir_dst=f"{args.model_dir}", + megatron_path=args.megatron_path, ) @@ -66,12 +84,12 @@ def prepare_cp(args: ScriptArgs): def _prepare_cp(args: ScriptArgs): U.rsync_simple( - path_src=f"/root/models/{args.model_name}_torch_dist", - path_dst=f"/root/local_data/{args.model_name}_torch_dist", + path_src=f"{args.model_dir}/{args.model_name}_torch_dist", + path_dst=f"{args.model_local_dir}/{args.model_name}_torch_dist", ) U.rsync_simple( - path_src=f"/root/models/{args.model_name}", - path_dst=f"/root/local_data/{args.model_name}", + path_src=f"{args.model_dir}/{args.model_name}", + path_dst=f"{args.model_local_dir}/{args.model_name}", ) @@ -80,12 +98,12 @@ def _prepare_cp(args: ScriptArgs): def train(args: ScriptArgs): print("running on {args.num_nodes} nodes") # ensure files are there is it was not synced before - _prepare_cp(args) + # _prepare_cp(args) - load_save_path = f"/root/shared_data/{args.run_id}/checkpoints" + load_save_path = f"{args.save_dir}/{args.run_id}/checkpoints" ckpt_args = ( - f"--hf-checkpoint /root/local_data/{args.model_name} " - f"--ref-load /root/local_data/{args.model_name}_torch_dist " + f"--hf-checkpoint {args.model_local_dir}/{args.model_name} " + f"--ref-load {args.model_local_dir}/{args.model_name}_torch_dist " f"--load {load_save_path} " f"--save {load_save_path} " "--save-interval 20 " @@ -120,24 +138,24 @@ def train(args: ScriptArgs): match args.task: case "dapo_aime": rollout_args += ( - "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl " "--input-key prompt " f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} " ) eval_args += ( - "--eval-prompt-data aime /root/datasets/aime-2024/aime-2024.jsonl " + f"--eval-prompt-data aime {args.data_dir}/aime-2024/aime-2024.jsonl " "--n-samples-per-eval-prompt 8 " "--eval-max-response-len 8192 " ) case "gsm8k": rollout_args += ( - "--prompt-data /root/datasets/gsm8k/train.parquet " + f"--prompt-data {args.data_dir}/gsm8k/train.parquet " "--input-key messages " # Deliberately make it very short for this easy task "--rollout-max-response-len 256 " ) eval_args += ( - "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet " + f"--eval-prompt-data gsm8k {args.data_dir}/gsm8k/test.parquet " "--n-samples-per-eval-prompt 1 " "--eval-max-response-len 256 " ) @@ -290,6 +308,7 @@ def train(args: ScriptArgs): num_gpus_per_node=args.num_gpus_per_node, megatron_model_type=args.megatron_model_type, extra_env_vars={**sglang_extra_env_vars}, + megatron_path=args.megatron_path, ) From cab9686fa801a5f9455392a29ba390f84bb226ab Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 20:01:19 -0800 Subject: [PATCH 20/30] fix --- .../run-qwen3-4b-mis.sh | 20 +++++++++---------- miles/backends/training_utils/loss.py | 1 - 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh b/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh index a130caa58f5..300e8ac75b1 100644 --- a/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh +++ b/examples/train_infer_mismatch_helper/run-qwen3-4b-mis.sh @@ -24,37 +24,37 @@ fi echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "/host_home/primary_synced/miles/scripts/models/qwen3-4B.sh" +source "/root/miles/scripts/models/qwen3-4B.sh" CKPT_ARGS=( - --hf-checkpoint /host_home/models/Qwen3-4B + --hf-checkpoint /root/Qwen3-4B #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/models/Qwen3-4B_torch_dist + --ref-load /root/Qwen3-4B_torch_dist # --load /root/Qwen3-4B_miles/ - --save /root/models/Qwen3-4B_miles/ + --save /root/Qwen3-4B_miles/ --save-interval 200 ) ROLLOUT_ARGS=( - --prompt-data /host_home/data/dapo-math-17k/dapo-math-17k.jsonl + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl --input-key prompt --label-key label --apply-chat-template --rollout-shuffle --rm-type deepscaler --num-rollout 3000 - --rollout-batch-size 8 + --rollout-batch-size 32 --n-samples-per-prompt 8 --rollout-max-response-len 8192 --rollout-temperature 1 - --global-batch-size 64 + --global-batch-size 256 --balance-data ) EVAL_ARGS=( # --eval-interval 20 - --eval-prompt-data aime /host_home/data/aime-2024/aime-2024.jsonl + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl --n-samples-per-eval-prompt 1 --eval-max-response-len 16384 --eval-top-p 1 @@ -127,7 +127,7 @@ CUSTOM_ARGS=( # launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 # Build the runtime environment JSON with proper variable substitution RUNTIME_ENV_JSON="{ @@ -142,7 +142,7 @@ ray job submit --address="http://127.0.0.1:8265" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ -- python3 train.py \ --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ + --actor-num-gpus-per-node 8 \ --colocate \ ${MODEL_ARGS[@]} \ ${CKPT_ARGS[@]} \ diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index 37e81eb533a..abc790761d0 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -87,7 +87,6 @@ def get_responses( tokens_chunk = tokens[-response_length:] else: # TODO: this is super ugly... do better abstraction. - max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp( total_length, response_length, parallel_state, qkv_format, max_seq_len ) From 8d51fe0bbb8f34d83174c4c4365d8cdfa967dfed Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sun, 28 Dec 2025 20:02:05 -0800 Subject: [PATCH 21/30] rm unused script --- scripts/run_deepseek_v3.2_5layer.py | 209 ---------------------------- 1 file changed, 209 deletions(-) delete mode 100644 scripts/run_deepseek_v3.2_5layer.py diff --git a/scripts/run_deepseek_v3.2_5layer.py b/scripts/run_deepseek_v3.2_5layer.py deleted file mode 100644 index 1643087e79e..00000000000 --- a/scripts/run_deepseek_v3.2_5layer.py +++ /dev/null @@ -1,209 +0,0 @@ -import re -from dataclasses import dataclass -from typing import Literal - -import typer - -import miles.utils.external_utils.command_utils as U - -app = typer.Typer() - - -@dataclass -class ScriptArgs(U.ExecuteTrainConfig): - run_id: str = U.create_run_id() - hf_checkpoint: str = "/root/.cache/dsv32-ckpt/DeepSeek-V3.2-5layer" - torch_dist_checkpoint: str = "/root/DeepSeek-V3.2-5layer_torch_dist" - num_gpus_per_node: int = 4 - enable_eval: bool = False - enable_deepep: bool = False - extra_args: str = "" - task: Literal["dapo_aime", "gsm8k"] = "dapo_aime" - mode: Literal["normal", "debug_minimal"] = "debug_minimal" - - -@app.command() -@U.dataclass_cli -def train(args: ScriptArgs): - load_save_path = f"/root/shared_data/{args.run_id}/checkpoints" - ckpt_args = ( - f"--hf-checkpoint {args.hf_checkpoint} " - f"--ref-load {args.torch_dist_checkpoint} " - f"--load {load_save_path} " - f"--save {load_save_path} " - "--save-interval 20 " - "--save-retain-interval 20 " - ) - - rollout_args = ( - "--label-key label " - "--apply-chat-template " - "--rollout-shuffle " - "--rm-type math " - "--num-rollout 3000 " - "--rollout-batch-size 128 " - "--n-samples-per-prompt 8 " - "--rollout-temperature 0.8 " - "--num-steps-per-rollout 4 " - "--balance-data " - ) - - if args.mode != "debug_minimal": - rollout_args += ( - "--over-sampling-batch-size 256 " - "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std " - ) - - eval_args = "" - if (args.mode != "debug_minimal") and args.enable_eval: - eval_args += "--eval-interval 20 " "--eval-top-p 0.7 " - - match args.task: - case "dapo_aime": - rollout_args += ( - "--prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl " - "--input-key prompt " - f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} " - ) - eval_args += ( - "--eval-prompt-data aime /root/aime-2024/aime-2024.jsonl " - "--n-samples-per-eval-prompt 8 " - "--eval-max-response-len 32768 " - ) - case "gsm8k": - rollout_args += ( - "--prompt-data /root/gsm8k/train.parquet " - "--input-key messages " - "--rollout-max-response-len 256 " - ) - eval_args += ( - "--eval-prompt-data gsm8k /root/gsm8k/test.parquet " - "--n-samples-per-eval-prompt 1 " - "--eval-max-response-len 256 " - ) - - if args.num_nodes <= 2: - perf_args = ( - "--tensor-model-parallel-size 1 " - "--sequence-parallel " - "--pipeline-model-parallel-size 1 " - "--context-parallel-size 8 " - f"--expert-model-parallel-size {args.num_gpus_per_node} " - "--expert-tensor-parallel-size 1 " - ) - elif args.num_nodes <= 4: - perf_args = ( - "--tensor-model-parallel-size 4 " - "--sequence-parallel " - "--pipeline-model-parallel-size 1 " - "--context-parallel-size 8 " - "--expert-model-parallel-size 8 " - "--expert-tensor-parallel-size 1 " - ) - else: - perf_args = ( - "--tensor-model-parallel-size 4 " - "--sequence-parallel " - "--pipeline-model-parallel-size 1 " - "--context-parallel-size 8 " - "--expert-model-parallel-size 16 " - "--expert-tensor-parallel-size 1 " - ) - perf_args += ( - "--recompute-granularity full " - "--recompute-method uniform " - "--recompute-num-layers 1 " - "--use-dynamic-batch-size " - "--max-tokens-per-gpu 2048 " - ) - - grpo_args = ( - "--advantage-estimator grpo " - "--kl-loss-coef 0.00 " - "--kl-loss-type low_var_kl " - "--entropy-coef 0.00 " - "--eps-clip 0.2 " - "--eps-clip-high 0.28 " - ) - - optimizer_args = ( - "--optimizer adam " - "--lr 1e-6 " - "--lr-decay-style constant " - "--weight-decay 0.1 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - ) - - sglang_decode_max_bs = 256 - sglang_world_size = args.num_gpus_per_node if args.num_nodes <= 4 else 64 - sglang_attn_dp_size = 1 if args.num_nodes <= 4 else 8 - sglang_attn_tp_size = sglang_world_size // sglang_attn_dp_size - sglang_args = ( - f"--rollout-num-gpus-per-engine {sglang_world_size} " - "--sglang-mem-fraction-static 0.7 " - # f"--sglang-tp-size {sglang_world_size} " - f"--sglang-tp-size 1 " - f"--sglang-ep-size {sglang_world_size} " - "--sglang-enable-dp-attention " - f"--sglang-dp-size {sglang_attn_dp_size} " - "--sglang-moe-dense-tp-size 1 " - "--sglang-enable-dp-lm-head " - "--sglang-server-concurrency 1024 " - f"--sglang-max-running-requests {sglang_world_size * sglang_decode_max_bs // sglang_attn_tp_size} " - f"--sglang-chunked-prefill-size {sglang_world_size * sglang_decode_max_bs} " - f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} " - ) - if args.enable_deepep: - sglang_args += ( - "--sglang-moe-a2a-backend deepep " - "--sglang-deepep-mode low_latency " - ) - sglang_extra_env_vars = {} - if args.enable_deepep: - sglang_extra_env_vars["SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = f"{sglang_decode_max_bs}" - - misc_args = ( - "--attention-dropout 0.0 " - "--hidden-dropout 0.0 " - "--accumulate-allreduce-grads-in-fp32 " - "--attention-softmax-in-fp32 " - "--attention-backend auto " - f"--update-weight-buffer-size {4 * 1024 ** 3} " - f"--actor-num-nodes {args.num_nodes} " - f"--actor-num-gpus-per-node {args.num_gpus_per_node} " - f"--num-gpus-per-node {args.num_gpus_per_node} " - "--colocate " - "--use-fault-tolerance " - f"--dump-details /root/shared_data/{args.run_id}/dump_details " - "--disable-weights-backuper " - "--model-name deepseekv32 " - "--train-memory-margin-bytes 1073741824 " - ) - - train_args = ( - f"{ckpt_args} " - f"{rollout_args} " - f"{optimizer_args} " - f"{grpo_args} " - f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} " - f"{perf_args} " - f"{eval_args} " - f"{sglang_args} " - f"{misc_args} " - f"{args.extra_args} " - ) - - U.execute_train( - train_args=train_args, - train_script="train.py", - config=args, - num_gpus_per_node=args.num_gpus_per_node, - megatron_model_type="deepseek-v32-5layer", - extra_env_vars={**sglang_extra_env_vars}, - ) - - -if __name__ == "__main__": - app() - From f7beab4f93fe5fdea1c7e8bfeca1eea9d963354e Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 29 Dec 2025 00:23:32 -0800 Subject: [PATCH 22/30] fix --- .../megatron_utils/megatron_to_hf/processors/quantizer_fp8.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index 54bb1e67646..87cf24992d7 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -42,7 +42,7 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio # TODO: find a clearer way. if converted_name.endswith("_scale"): continue - if_use_ue8m0_in_moe = True if args.sglang_moe_a2a_backend == "deepep" else False + if_use_ue8m0_in_moe = True if args.sglang_moe_runner_backend == "deep_gemm" else False quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size, if_use_ue8m0_in_moe=if_use_ue8m0_in_moe)) return quantize_named_params From dd6870636290c5c00a156c21512fa755e8019eee Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 29 Dec 2025 13:42:29 -0800 Subject: [PATCH 23/30] add docs --- docker/deepseekv32/README.md | 41 ++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 docker/deepseekv32/README.md diff --git a/docker/deepseekv32/README.md b/docker/deepseekv32/README.md new file mode 100644 index 00000000000..1ca43f1c2bd --- /dev/null +++ b/docker/deepseekv32/README.md @@ -0,0 +1,41 @@ +## Usage + +### Docker +```bash +docker pull yueming11/miles:dsv32-dev + +docker run --gpus all --ipc=host --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --name miles_dsv32 yueming11/miles:dsv32-dev /bin/zsh + +git clone https://github.com/radixark/miles.git +git checkout dsv32 +cd dsv32 +pip install -e . + +# if shows Megatron does not support numpy 2.x +pip install numpy==1.26.4 +``` + +### Quick test with 5 layer model +#### model download + +``` +hf download Pinaster/DeepSeek-V3.2-5layer /root/models/DeepSeek-V3.2-5layer +``` + +#### Prepare model for training +Note: need to change the paths, for all commands below see `scripts/run_deepseek_v32.py` for details + +Step 1. download dataset & convert fp8 hf checkpoint to bf16 with one node +``` +python scripts/run_deepseek_v32.py prepare-single --model-name DeepSeek-V3.2-5layer --megatron-model-type deepseek-v32-5layer +``` + +Step 2. convert hf checkpoint to megatron checkpoint with multiple nodes +``` +python scripts/run_deepseek_v32.py prepare-spmd --model-name DeepSeek-V3.2-5layer --megatron-model-type deepseek-v32-5layer +``` + +#### Launch training +``` +python scripts/run_deepseek_v32.py train --model-name DeepSeek-V3.2-5layer --megatron-model-type deepseek-v32-5layer +``` \ No newline at end of file From f16e095957910bd6b353a5a8216870d911ef82f0 Mon Sep 17 00:00:00 2001 From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com> Date: Tue, 6 Jan 2026 22:05:05 -0600 Subject: [PATCH 24/30] Fix torch native CP attention backend for DSA (#406) --- docker/deepseekv32/megatron.patch | 245 ++++++++++++++++++++++++++---- 1 file changed, 218 insertions(+), 27 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index e6a8a34a561..3914600938a 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -1,8 +1,77 @@ diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py -index 89659a1d7..38efa896c 100644 +index 89659a1d7..77f1beb87 100644 --- a/megatron/core/transformer/dot_product_attention_context_parallel.py +++ b/megatron/core/transformer/dot_product_attention_context_parallel.py -@@ -132,10 +132,10 @@ class AllGatherComm: +@@ -6,6 +6,7 @@ + + import torch + from torch.nn import functional as F ++import torch.distributed as dist + + try: + import einops +@@ -53,15 +54,19 @@ def eager_attn_fwd(q, k, v, attn_bias, sinks, scale, dropout): + + + @torch.no_grad +-def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs, grad_output): ++def eager_attn_bwd(q, kv, attn_bias, sinks, scale, dim_short, dropout, attn_output, probs, grad_output): + """Backward pass for eager attention""" + + # Rearrange query, key, value to (b, h, s, d) + b, sq, h, d = q.shape +- sk = k.shape[1] ++ _, sk, _, _ = kv.shape ++ k = kv ++ v = kv[:,:,:,:dim_short] ++ q_tail = q[:,:,:,dim_short:] ++ _q_tail_T = einops.rearrange(q_tail, 'b s h d -> b h d s').contiguous() + _q_T = einops.rearrange(q, 'b s h d -> b h d s') + _k_T = einops.rearrange(k, 'b s h d -> b h s d') +- _v_T = einops.rearrange(v, ' b s h d -> b h d s') ++ _v_T = einops.rearrange(v, 'b s h d -> b h d s') + + # Backward pass for score @ value + if sinks is None: +@@ -70,9 +75,9 @@ def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs + attn_w = probs[..., :-1] # Drop the sink + grad_output = einops.rearrange(grad_output, 'b s h d -> b h s d') + attn_w_T = einops.rearrange(attn_w, ' b h sq sk -> b h sk sq') +- grad__v = torch.matmul(attn_w_T, grad_output) +- grad_attn_w = torch.matmul(grad_output, _v_T) +- ++ grad__v = torch.matmul(attn_w_T, grad_output).contiguous() # b h sk d ++ grad_attn_w = torch.matmul(grad_output, _v_T).contiguous() # b h s d || b h d sk -> b h s sk ++ + # Backward pass for softmax + if sinks is None: + grad_probs = grad_attn_w +@@ -95,15 +100,18 @@ def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs + + # Backward pass for q @ K^T + grad_attn_w *= scale +- grad__q = torch.matmul(grad_attn_w, _k_T) +- grad__k = torch.matmul(_q_T, grad_attn_w) ++ grad__q = torch.matmul(grad_attn_w, _k_T).contiguous() ++ grad__k = torch.matmul(_q_T, grad_attn_w).contiguous() # b h d sk ++ ++ grad__k_T = grad__k.transpose(2, 3).contiguous() # b h sk d ++ grad__kv = torch.zeros((b, h, sk, d), device=q.device, dtype=q.dtype) # b h sk d ++ grad__kv[:,:,:,:dim_short] = grad__v + grad__k_T[:,:,:,:dim_short] ++ grad__kv[:,:,:,dim_short:] = torch.matmul(_q_tail_T, grad_attn_w).contiguous().transpose(2, 3).contiguous() # b h sk d + + # Rearrange grads to (b, s, h, d) +- grad_v = einops.rearrange(grad__v, 'b h s d -> b s h d') +- grad_k = einops.rearrange(grad__k, 'b h d s -> b s h d') ++ grad__kv = grad__kv.transpose(1, 2).contiguous() + grad_q = einops.rearrange(grad__q, 'b h s d -> b s h d') +- return grad_q, grad_k, grad_v, grad_sinks +- ++ return grad_q, grad__kv, grad_sinks + + class AllGatherComm: + """All gather communication with async operations""" +@@ -132,10 +140,10 @@ class AllGatherComm: self.handles = [] @@ -15,7 +84,7 @@ index 89659a1d7..38efa896c 100644 zz_mask = attention_mask else: chunked = attention_mask.chunk(dim=3, chunks=cp_size * 2) -@@ -151,7 +151,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +@@ -151,7 +159,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): """Native attention function with context parallelism.""" @staticmethod @@ -24,7 +93,7 @@ index 89659a1d7..38efa896c 100644 '''Forward pass for the native attention function with context parallelism''' # Assert einops exists -@@ -171,12 +171,17 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +@@ -171,12 +179,17 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): probs = [] # Initialize KV buffers @@ -46,7 +115,7 @@ index 89659a1d7..38efa896c 100644 # All-gather first chunk of KV buffers k_0 = k[:, :, :heads_k_stride].contiguous() -@@ -186,7 +191,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +@@ -186,7 +199,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): # Prepare attention bias attn_bias = to_zz_mask_attn_bias( @@ -55,7 +124,18 @@ index 89659a1d7..38efa896c 100644 ) # Iterate over heads -@@ -226,6 +231,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +@@ -215,8 +228,9 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + + # Forward pass + out_i, probs_i = eager_attn_fwd( +- q_i, k_i, v_i, attn_bias, None, softmax_scale, attention_dropout ++ q_i, k_i, v_i, attn_bias.contiguous(), None, softmax_scale, attention_dropout + ) ++ + outs.append(out_i) + probs.append(probs_i) + +@@ -226,10 +240,13 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): # Save contexts for backward pass ctx.save_for_backward(q, k, v, attention_mask, *outs, *probs) @@ -63,42 +143,153 @@ index 89659a1d7..38efa896c 100644 ctx.dropout = attention_dropout ctx.scale = softmax_scale ctx.heads_k_stride = heads_k_stride # TODO make it configurable -@@ -252,12 +258,16 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): - comm = AllGatherComm(group=pg) + ctx.pg = pg ++ ctx.dim = q.shape[3] ++ ctx.dim_short = v.shape[3] + + return out + +@@ -238,13 +255,15 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + '''Backward pass for the native attention function with context parallelism''' + + # Initialize or resume constants and communication group +- q, k, v, attention_mask, *rest = ctx.saved_tensors ++ q, kv, _, attention_mask, *rest = ctx.saved_tensors ++ dim = ctx.dim ++ dim_short = ctx.dim_short + nheads = q.shape[2] +- nheads_k = k.shape[2] +- heads_k_stride = ctx.heads_k_stride +- assert nheads_k % heads_k_stride == 0 +- outs = rest[: nheads_k // heads_k_stride] +- probs = rest[nheads_k // heads_k_stride :] ++ nheads_kv = kv.shape[2] ++ heads_kv_stride = ctx.heads_k_stride ++ assert nheads_kv % heads_kv_stride == 0 ++ outs = rest[: nheads_kv // heads_kv_stride] ++ probs = rest[nheads_kv // heads_kv_stride :] + pg = ctx.pg + cp_size = 1 + if pg is not None: +@@ -253,30 +272,27 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): # Initialize KV buffers -- kv_buffer = torch.empty( + kv_buffer = torch.empty( - (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), -+ kv_buffer = [torch.empty( -+ (k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), - dtype=k.dtype, - device=k.device, -- ) -- kv_buffer_copy = torch.empty_like(kv_buffer) -+ ), torch.empty( -+ (v.shape[0] * cp_size, v.shape[1], heads_k_stride, v.shape[3]), -+ dtype=v.dtype, -+ device=v.device, -+ )] -+ kv_buffer_copy = [torch.empty_like(kv_buffer[0]), torch.empty_like(kv_buffer[1])] +- dtype=k.dtype, +- device=k.device, ++ (kv.shape[0] * cp_size, kv.shape[1], heads_kv_stride, kv.shape[3]), ++ dtype=kv.dtype, ++ device=kv.device, + ) + kv_buffer_copy = torch.empty_like(kv_buffer) # All-gather first chunk of KV buffers dq = [] -@@ -270,7 +280,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +- dk = [] +- dv = [] +- k_0 = k[:, :, :heads_k_stride].contiguous() +- v_0 = v[:, :, :heads_k_stride].contiguous() +- comm.all_gather(kv_buffer_copy[0], k_0) +- comm.all_gather(kv_buffer_copy[1], v_0) ++ dkv = [] ++ kv_0 = kv[:, :, :heads_kv_stride].contiguous() ++ comm.all_gather(kv_buffer_copy, kv_0) # Prepare attention bias attn_bias = to_zz_mask_attn_bias( - attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype -+ attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype, ctx.if_zz_mask ++ attention_mask, cp_size, nheads, nheads_kv, heads_kv_stride, q.device, q.dtype, ctx.if_zz_mask ) # Iterate over heads -@@ -339,4 +349,4 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +- for i in range(0, nheads_k, heads_k_stride): ++ for i in range(0, nheads_kv, heads_kv_stride): + # Slice query and output for this iteration +- q_slice = slice(i * nheads // nheads_k, (i + heads_k_stride) * nheads // nheads_k) ++ q_slice = slice(i * nheads // nheads_kv, (i + heads_kv_stride) * nheads // nheads_kv) + q_i = q[:, :, q_slice] + dout_i = dout[:, :, q_slice] + +@@ -285,58 +301,45 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer + + # All-gather the next portion of KV buffers if not the last iteration +- if i < nheads_k - heads_k_stride: +- kvsl = i + heads_k_stride +- kvsr = kvsl + heads_k_stride +- send_k = k[:, :, kvsl:kvsr].contiguous() +- send_v = v[:, :, kvsl:kvsr].contiguous() +- comm.all_gather(kv_buffer_copy[0], send_k) +- comm.all_gather(kv_buffer_copy[1], send_v) ++ if i < nheads_kv - heads_kv_stride: ++ kvsl = i + heads_kv_stride ++ kvsr = kvsl + heads_kv_stride ++ send_kv = kv[:, :, kvsl:kvsr].contiguous() ++ comm.all_gather(kv_buffer_copy, send_kv) + + # Prepare key, value for attention +- k_i = kv_buffer[0] +- v_i = kv_buffer[1] ++ kv_i = kv_buffer + + # Rearrange query, key, value to (b, s, h, d) + q_i = einops.rearrange(q_i, 's b h d -> b s h d') +- k_i = einops.rearrange(k_i, 's b h d -> b s h d') +- v_i = einops.rearrange(v_i, 's b h d -> b s h d') ++ kv_i = einops.rearrange(kv_i, 's b h d -> b s h d') + dout_i = einops.rearrange(dout_i, 's b h d -> b s h d') + + # Backward pass +- dq_i, _dk_i, _dv_i, _ = eager_attn_bwd( +- q_i, k_i, v_i, attn_bias, None, ctx.scale, ctx.dropout, outs[i], probs[i], dout_i ++ dq_i, _dkv_i, _ = eager_attn_bwd( ++ q_i, kv_i, attn_bias, None, ctx.scale, dim_short, ctx.dropout, outs[i], probs[i], dout_i + ) + + # Rearrange gradients to (s, b, h, d) + dq_i = einops.rearrange(dq_i, 'b s h d -> s b h d') +- _dk_i = einops.rearrange(_dk_i, 'b s h d -> s b h d') +- _dv_i = einops.rearrange(_dv_i, 'b s h d -> s b h d') ++ _dkv_i = einops.rearrange(_dkv_i, 'b s h d -> s b h d') ++ + if pg is None: +- dk_i = _dk_i +- dv_i = _dv_i ++ dkv_i = _dkv_i + else: + # Reduce-scatter gradients if CP > 1 +- dk_i = torch.zeros( +- (k_i.shape[1] // cp_size, k_i.shape[0], k_i.shape[2], k_i.shape[3]), +- device=k_i.device, +- dtype=k_i.dtype, +- ) +- dv_i = torch.zeros( +- (v_i.shape[1] // cp_size, v_i.shape[0], v_i.shape[2], v_i.shape[3]), +- device=v_i.device, +- dtype=v_i.dtype, ++ dkv_i = torch.zeros( ++ (kv_i.shape[1] // cp_size, kv_i.shape[0], kv_i.shape[2], kv_i.shape[3]), ++ device=kv_i.device, ++ dtype=kv_i.dtype, + ) +- torch.distributed.reduce_scatter_tensor(dk_i, _dk_i, group=pg) +- torch.distributed.reduce_scatter_tensor(dv_i, _dv_i, group=pg) ++ torch.distributed.reduce_scatter_tensor(dkv_i, _dkv_i, group=pg) + + # Collect gradients + dq.append(dq_i) +- dk.append(dk_i) +- dv.append(dv_i) ++ dkv.append(dkv_i) + + # Concatenate gradients and return dq = torch.cat(dq, dim=2) - dk = torch.cat(dk, dim=2) - dv = torch.cat(dv, dim=2) +- dk = torch.cat(dk, dim=2) +- dv = torch.cat(dv, dim=2) - return dq, dk, dv, None, None, None, None -+ return dq, dk, dv, None, None, None, None, None ++ dkv = torch.cat(dkv, dim=2) ++ return dq, dkv, dkv[:,:,:,:dim_short].detach().contiguous(), None, None, None, None, None diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index fc994490b..7bc9a485e 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py From e28d439e4ae6554de792844f455948d0dd4afc1c Mon Sep 17 00:00:00 2001 From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com> Date: Fri, 16 Jan 2026 14:21:48 +0800 Subject: [PATCH 25/30] tilelang kernel + matrix absorb in megatron (#461) --- docker/deepseekv32/megatron.patch | 1296 ++++++++++++++--- .../processors/quantizer_fp8.py | 12 +- .../megatron_utils/update_weight/common.py | 12 +- miles_plugins/mbridge/__init__.py | 14 +- miles_plugins/mbridge/deepseekv32.py | 68 +- scripts/run_deepseek_v32.py | 5 +- 6 files changed, 1139 insertions(+), 268 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index 3914600938a..ac7a1be3c1c 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -1,220 +1,324 @@ diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py -index 89659a1d7..77f1beb87 100644 +index 89659a1d7..1def27c69 100644 --- a/megatron/core/transformer/dot_product_attention_context_parallel.py +++ b/megatron/core/transformer/dot_product_attention_context_parallel.py -@@ -6,6 +6,7 @@ +@@ -3,9 +3,12 @@ + # Some of this code was adopted from https://github.com/zhuzilin/ring-flash-attention/ + # This source code is licensed under the MIT license found in the + # LICENSE file in the root directory of this source tree. ++# Kernel is adpoted from tilelang/examples/deepseek_v32 import torch - from torch.nn import functional as F +import torch.distributed as dist + from torch.nn import functional as F ++from .tilelang_kernel import sparse_mla_bwd, sparse_mla_fwd_interface try: import einops -@@ -53,15 +54,19 @@ def eager_attn_fwd(q, k, v, attn_bias, sinks, scale, dropout): +@@ -15,96 +18,6 @@ except ImportError: + HAVE_EINOPS = False - @torch.no_grad +-@torch.no_grad +-def eager_attn_fwd(q, k, v, attn_bias, sinks, scale, dropout): +- """Forward pass for eager attention""" +- +- # Rearrange query, key, value to (b, h, s, d) +- b, sq, h, d = q.shape +- sk = k.shape[1] +- _q = einops.rearrange(q, 'b s h d -> b h s d') +- _k = einops.rearrange(k, 'b s h d -> b h d s') +- _v = einops.rearrange(v, 'b s h d -> b h s d') +- +- # Compute attention weights +- attn_w = torch.matmul(_q, _k) * scale +- attn_w = attn_w + attn_bias +- +- # Add sinks to attention weights +- if sinks is None: +- logits = attn_w +- else: +- _sinks = sinks.reshape(1, h, 1, 1).expand(b, -1, sq, 1) +- logits = torch.cat([attn_w, _sinks], dim=-1) +- +- # Compute attention scores +- probs = F.softmax(logits, dim=-1, dtype=logits.dtype) +- if sinks is None: +- attn_w = probs +- else: +- attn_w = probs[..., :-1] # Drop the sink +- +- # Compute attention output +- attn_output = torch.matmul(attn_w, _v) +- attn_output = einops.rearrange(attn_output, 'b h s d -> b s h d') +- attn_output = attn_output.contiguous() +- +- return attn_output, probs +- +- +-@torch.no_grad -def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs, grad_output): -+def eager_attn_bwd(q, kv, attn_bias, sinks, scale, dim_short, dropout, attn_output, probs, grad_output): - """Backward pass for eager attention""" - - # Rearrange query, key, value to (b, h, s, d) - b, sq, h, d = q.shape +- """Backward pass for eager attention""" +- +- # Rearrange query, key, value to (b, h, s, d) +- b, sq, h, d = q.shape - sk = k.shape[1] -+ _, sk, _, _ = kv.shape -+ k = kv -+ v = kv[:,:,:,:dim_short] -+ q_tail = q[:,:,:,dim_short:] -+ _q_tail_T = einops.rearrange(q_tail, 'b s h d -> b h d s').contiguous() - _q_T = einops.rearrange(q, 'b s h d -> b h d s') - _k_T = einops.rearrange(k, 'b s h d -> b h s d') +- _q_T = einops.rearrange(q, 'b s h d -> b h d s') +- _k_T = einops.rearrange(k, 'b s h d -> b h s d') - _v_T = einops.rearrange(v, ' b s h d -> b h d s') -+ _v_T = einops.rearrange(v, 'b s h d -> b h d s') - - # Backward pass for score @ value - if sinks is None: -@@ -70,9 +75,9 @@ def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs - attn_w = probs[..., :-1] # Drop the sink - grad_output = einops.rearrange(grad_output, 'b s h d -> b h s d') - attn_w_T = einops.rearrange(attn_w, ' b h sq sk -> b h sk sq') +- +- # Backward pass for score @ value +- if sinks is None: +- attn_w = probs +- else: +- attn_w = probs[..., :-1] # Drop the sink +- grad_output = einops.rearrange(grad_output, 'b s h d -> b h s d') +- attn_w_T = einops.rearrange(attn_w, ' b h sq sk -> b h sk sq') - grad__v = torch.matmul(attn_w_T, grad_output) - grad_attn_w = torch.matmul(grad_output, _v_T) - -+ grad__v = torch.matmul(attn_w_T, grad_output).contiguous() # b h sk d -+ grad_attn_w = torch.matmul(grad_output, _v_T).contiguous() # b h s d || b h d sk -> b h s sk -+ - # Backward pass for softmax - if sinks is None: - grad_probs = grad_attn_w -@@ -95,15 +100,18 @@ def eager_attn_bwd(q, k, v, attn_bias, sinks, scale, dropout, attn_output, probs - - # Backward pass for q @ K^T - grad_attn_w *= scale +- # Backward pass for softmax +- if sinks is None: +- grad_probs = grad_attn_w +- else: +- dummy = torch.zeros((b, h, sq, 1), device=q.device, dtype=q.dtype) +- grad_probs = torch.cat([grad_attn_w, dummy], dim=3) +- del grad_attn_w +- grad_logits = torch._softmax_backward_data( +- grad_probs, probs, -1, probs.dtype +- ) # [b, h, sq, sk+1] +- +- # Backward pass for adding sinks +- if sinks is None: +- grad_sinks = None +- grad_attn_w = grad_logits +- else: +- grad__sinks = grad_logits[:, :, :, -1] # [b, h, sq] +- grad_sinks = einops.rearrange(grad__sinks, 'b h s -> h (b s)').sum(-1) +- grad_attn_w = grad_logits[:, :, :, :-1].contiguous() # [b, h, sq, sk] +- +- # Backward pass for q @ K^T +- grad_attn_w *= scale - grad__q = torch.matmul(grad_attn_w, _k_T) - grad__k = torch.matmul(_q_T, grad_attn_w) -+ grad__q = torch.matmul(grad_attn_w, _k_T).contiguous() -+ grad__k = torch.matmul(_q_T, grad_attn_w).contiguous() # b h d sk -+ -+ grad__k_T = grad__k.transpose(2, 3).contiguous() # b h sk d -+ grad__kv = torch.zeros((b, h, sk, d), device=q.device, dtype=q.dtype) # b h sk d -+ grad__kv[:,:,:,:dim_short] = grad__v + grad__k_T[:,:,:,:dim_short] -+ grad__kv[:,:,:,dim_short:] = torch.matmul(_q_tail_T, grad_attn_w).contiguous().transpose(2, 3).contiguous() # b h sk d - - # Rearrange grads to (b, s, h, d) +- +- # Rearrange grads to (b, s, h, d) - grad_v = einops.rearrange(grad__v, 'b h s d -> b s h d') - grad_k = einops.rearrange(grad__k, 'b h d s -> b s h d') -+ grad__kv = grad__kv.transpose(1, 2).contiguous() - grad_q = einops.rearrange(grad__q, 'b h s d -> b s h d') +- grad_q = einops.rearrange(grad__q, 'b h s d -> b s h d') - return grad_q, grad_k, grad_v, grad_sinks - -+ return grad_q, grad__kv, grad_sinks - +- class AllGatherComm: """All gather communication with async operations""" -@@ -132,10 +140,10 @@ class AllGatherComm: - self.handles = [] +@@ -131,212 +44,146 @@ class AllGatherComm: + handle.wait() + self.handles = [] +- -def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stride, device, dtype): -+def to_zz_mask_attn_bias(attention_mask, cp_size, nheads, nheads_k, heads_k_stride, device, dtype, if_zz_mask=False): - '''Convert the attention mask to the attention bias''' - +- '''Convert the attention mask to the attention bias''' +- - if cp_size == 1: -+ if cp_size == 1 or if_zz_mask: - zz_mask = attention_mask - else: - chunked = attention_mask.chunk(dim=3, chunks=cp_size * 2) -@@ -151,7 +159,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): +- zz_mask = attention_mask +- else: +- chunked = attention_mask.chunk(dim=3, chunks=cp_size * 2) +- zz_mask = [_x for _p in zip(chunked[:cp_size], reversed(chunked[cp_size:])) for _x in _p] +- zz_mask = torch.cat(zz_mask, dim=3) +- attn_bias = torch.zeros(zz_mask.shape, device=device, dtype=dtype) +- attn_bias.masked_fill_(zz_mask, float('-inf')) +- attn_bias = attn_bias.expand(-1, heads_k_stride * (nheads // nheads_k), -1, -1) +- return attn_bias +- +- + class AttentionFuncionWithContextParallel(torch.autograd.Function): """Native attention function with context parallelism.""" ++ # q: [seq_len_shard, batch, nheads, dim] ++ # k: [seq_len_kv_shard, batch, 1, dim] ++ # v: [seq_len_kv_shard, batch, 1, dim_v] ++ # indices: [batch, 1, seq_len, topk] ++ # masks: [batch, 1, seq_len, seq_len_kv] @staticmethod - def forward(ctx, q, k, v, attention_mask, attention_dropout, softmax_scale, pg): -+ def forward(ctx, q, k, v, attention_mask, attention_dropout, softmax_scale, pg, if_zz_mask=False): ++ def forward(ctx, q, k, dim_v, indices, masks, attention_dropout, softmax_scale, pg): '''Forward pass for the native attention function with context parallelism''' - # Assert einops exists -@@ -171,12 +179,17 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): - probs = [] +- # Assert einops exists + if not HAVE_EINOPS: + raise ImportError("einops is required by the attention CP but cannot be imported.") - # Initialize KV buffers +- # Initialize communication group and constants + cp_size = 1 + if pg is not None: + cp_size = torch.distributed.get_world_size(pg) + comm = AllGatherComm(group=pg) +- nheads = q.shape[2] +- nheads_k = k.shape[2] +- heads_k_stride = 1 +- assert nheads % nheads_k == 0 and nheads_k % heads_k_stride == 0 +- outs = [] +- probs = [] +- +- # Initialize KV buffers - kv_buffer = torch.empty( - (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), -+ # seperate KV buffer for MLA -+ kv_buffer = [torch.empty( -+ (k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), ++ s, b, heads, dim = q.shape ++ skv, _, kv_groups, _ = k.shape ++ ++ k_buffer = torch.empty( ++ (k.shape[0] * cp_size, k.shape[1], 1, k.shape[3]), dtype=k.dtype, device=k.device, -- ) + ) - kv_buffer_copy = torch.empty_like(kv_buffer) -+ ), torch.empty( -+ (v.shape[0] * cp_size, v.shape[1], heads_k_stride, v.shape[3]), -+ dtype=v.dtype, -+ device=v.device, -+ )] -+ kv_buffer_copy = [torch.empty_like(kv_buffer[0]), torch.empty_like(kv_buffer[1])] - - # All-gather first chunk of KV buffers - k_0 = k[:, :, :heads_k_stride].contiguous() -@@ -186,7 +199,7 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): - - # Prepare attention bias - attn_bias = to_zz_mask_attn_bias( +- +- # All-gather first chunk of KV buffers +- k_0 = k[:, :, :heads_k_stride].contiguous() +- v_0 = v[:, :, :heads_k_stride].contiguous() +- comm.all_gather(kv_buffer_copy[0], k_0) +- comm.all_gather(kv_buffer_copy[1], v_0) +- +- # Prepare attention bias +- attn_bias = to_zz_mask_attn_bias( - attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype -+ attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype, if_zz_mask - ) - - # Iterate over heads -@@ -215,8 +228,9 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): - - # Forward pass - out_i, probs_i = eager_attn_fwd( +- ) +- +- # Iterate over heads +- for i in range(0, nheads_k, heads_k_stride): +- # Wait for previous all-gather to complete +- comm.wait() +- kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer +- # All-gather the next portion of KV buffers if not the last iteration +- if i < nheads_k - heads_k_stride: +- kvsl = i + heads_k_stride +- kvsr = kvsl + heads_k_stride +- send_k = k[:, :, kvsl:kvsr].contiguous() +- send_v = v[:, :, kvsl:kvsr].contiguous() +- comm.all_gather(kv_buffer_copy[0], send_k) +- comm.all_gather(kv_buffer_copy[1], send_v) +- +- # Prepare query, key, value for attention +- q_i = q[:, :, i * nheads // nheads_k : (i + heads_k_stride) * nheads // nheads_k] +- k_i = kv_buffer[0] +- v_i = kv_buffer[1] +- +- # Rearrange query, key, value to (b, s, h, d) +- q_i = einops.rearrange(q_i, 's b h d -> b s h d') +- k_i = einops.rearrange(k_i, 's b h d -> b s h d') +- v_i = einops.rearrange(v_i, 's b h d -> b s h d') +- +- # Forward pass +- out_i, probs_i = eager_attn_fwd( - q_i, k_i, v_i, attn_bias, None, softmax_scale, attention_dropout -+ q_i, k_i, v_i, attn_bias.contiguous(), None, softmax_scale, attention_dropout - ) +- ) +- outs.append(out_i) +- probs.append(probs_i) +- +- # Concatenate outputs and rearrange to (s, b, h, d) +- out = torch.cat(outs, dim=2) +- out = einops.rearrange(out, 'b s h d -> s b h d') +- +- # Save contexts for backward pass +- ctx.save_for_backward(q, k, v, attention_mask, *outs, *probs) ++ comm.all_gather(k_buffer, k) ++ comm.wait() + - outs.append(out_i) - probs.append(probs_i) - -@@ -226,10 +240,13 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): - - # Save contexts for backward pass - ctx.save_for_backward(q, k, v, attention_mask, *outs, *probs) -+ ctx.if_zz_mask = if_zz_mask ++ zz_indices = indices.transpose(1, 2) ++ zz_masks = masks.transpose(1, 2) ++ ++ q_i = q ++ k_i = k_buffer ++ ++ s_, b_, h_, d_ = q_i.shape ++ q_i = einops.rearrange(q_i, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = k_i.shape ++ k_i = einops.rearrange(k_i, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ zz_indices_i = zz_indices ++ b_, s_, g_, topk_ = zz_indices_i.shape ++ zz_indices_i = zz_indices_i.flatten().view(b_, s_, g_, topk_) ++ zz_masks_i = zz_masks ++ b_, s_, g_, skv_ = zz_masks_i.shape ++ zz_masks_i = zz_masks_i.flatten().view(b_, s_, g_, skv_) ++ ++ out_i, lse_i = sparse_mla_fwd_interface(q_i.contiguous(), k_i, zz_indices_i, zz_masks_i, dim_v, sm_scale = softmax_scale) ++ ++ # out: [B, seq_len_shard, h, dim] -> [seq_len, B, h, dim] ++ out_i = einops.rearrange(out_i, 'b s h d -> s b h d') ++ ++ # outs: [[B, seq_len_shard, nheads // kv_group, dim], ...., [B, seq_len_shard, nheads // kv_group, dim]], repeat kv_group // heads_kv_stride times ++ # lses: [[B, seq_len_shard, heads_kv_stride], ...., [B, seq_len_shard, heads_kv_stride]], repeat kv_group // heads_kv_stride times ++ ctx.save_for_backward(q, k, indices, masks, out_i, lse_i) ctx.dropout = attention_dropout - ctx.scale = softmax_scale - ctx.heads_k_stride = heads_k_stride # TODO make it configurable +- ctx.scale = softmax_scale +- ctx.heads_k_stride = heads_k_stride # TODO make it configurable ++ ctx.softmax_scale = softmax_scale ++ ctx.dim_v = dim_v ctx.pg = pg -+ ctx.dim = q.shape[3] -+ ctx.dim_short = v.shape[3] - return out +- return out ++ return out_i -@@ -238,13 +255,15 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + @staticmethod + def backward(ctx, dout): '''Backward pass for the native attention function with context parallelism''' - # Initialize or resume constants and communication group +- # Initialize or resume constants and communication group - q, k, v, attention_mask, *rest = ctx.saved_tensors -+ q, kv, _, attention_mask, *rest = ctx.saved_tensors -+ dim = ctx.dim -+ dim_short = ctx.dim_short - nheads = q.shape[2] +- nheads = q.shape[2] - nheads_k = k.shape[2] - heads_k_stride = ctx.heads_k_stride - assert nheads_k % heads_k_stride == 0 - outs = rest[: nheads_k // heads_k_stride] - probs = rest[nheads_k // heads_k_stride :] -+ nheads_kv = kv.shape[2] -+ heads_kv_stride = ctx.heads_k_stride -+ assert nheads_kv % heads_kv_stride == 0 -+ outs = rest[: nheads_kv // heads_kv_stride] -+ probs = rest[nheads_kv // heads_kv_stride :] ++ q, k, indices, masks, out, lse = ctx.saved_tensors ++ s, b, heads, dim = q.shape ++ dim_v = ctx.dim_v ++ softmax_scale = ctx.softmax_scale ++ pg = ctx.pg cp_size = 1 if pg is not None: -@@ -253,30 +272,27 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): + cp_size = torch.distributed.get_world_size(pg) + comm = AllGatherComm(group=pg) - # Initialize KV buffers - kv_buffer = torch.empty( +- # Initialize KV buffers +- kv_buffer = torch.empty( - (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), -- dtype=k.dtype, -- device=k.device, -+ (kv.shape[0] * cp_size, kv.shape[1], heads_kv_stride, kv.shape[3]), -+ dtype=kv.dtype, -+ device=kv.device, ++ k_buffer = torch.empty( ++ (k.shape[0] * cp_size, k.shape[1], 1, k.shape[3]), + dtype=k.dtype, + device=k.device, ) - kv_buffer_copy = torch.empty_like(kv_buffer) - - # All-gather first chunk of KV buffers - dq = [] +- kv_buffer_copy = torch.empty_like(kv_buffer) +- +- # All-gather first chunk of KV buffers +- dq = [] - dk = [] - dv = [] - k_0 = k[:, :, :heads_k_stride].contiguous() - v_0 = v[:, :, :heads_k_stride].contiguous() - comm.all_gather(kv_buffer_copy[0], k_0) - comm.all_gather(kv_buffer_copy[1], v_0) -+ dkv = [] -+ kv_0 = kv[:, :, :heads_kv_stride].contiguous() -+ comm.all_gather(kv_buffer_copy, kv_0) - - # Prepare attention bias - attn_bias = to_zz_mask_attn_bias( +- +- # Prepare attention bias +- attn_bias = to_zz_mask_attn_bias( - attention_mask, cp_size, nheads, nheads_k, heads_k_stride, q.device, q.dtype -+ attention_mask, cp_size, nheads, nheads_kv, heads_kv_stride, q.device, q.dtype, ctx.if_zz_mask - ) +- ) - # Iterate over heads +- # Iterate over heads - for i in range(0, nheads_k, heads_k_stride): -+ for i in range(0, nheads_kv, heads_kv_stride): - # Slice query and output for this iteration +- # Slice query and output for this iteration - q_slice = slice(i * nheads // nheads_k, (i + heads_k_stride) * nheads // nheads_k) -+ q_slice = slice(i * nheads // nheads_kv, (i + heads_kv_stride) * nheads // nheads_kv) - q_i = q[:, :, q_slice] - dout_i = dout[:, :, q_slice] - -@@ -285,58 +301,45 @@ class AttentionFuncionWithContextParallel(torch.autograd.Function): - kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer - - # All-gather the next portion of KV buffers if not the last iteration +- q_i = q[:, :, q_slice] +- dout_i = dout[:, :, q_slice] +- +- # Wait for previous all-gather to complete +- comm.wait() +- kv_buffer, kv_buffer_copy = kv_buffer_copy, kv_buffer +- +- # All-gather the next portion of KV buffers if not the last iteration - if i < nheads_k - heads_k_stride: - kvsl = i + heads_k_stride - kvsr = kvsl + heads_k_stride @@ -222,76 +326,101 @@ index 89659a1d7..77f1beb87 100644 - send_v = v[:, :, kvsl:kvsr].contiguous() - comm.all_gather(kv_buffer_copy[0], send_k) - comm.all_gather(kv_buffer_copy[1], send_v) -+ if i < nheads_kv - heads_kv_stride: -+ kvsl = i + heads_kv_stride -+ kvsr = kvsl + heads_kv_stride -+ send_kv = kv[:, :, kvsl:kvsr].contiguous() -+ comm.all_gather(kv_buffer_copy, send_kv) - - # Prepare key, value for attention +- +- # Prepare key, value for attention - k_i = kv_buffer[0] - v_i = kv_buffer[1] -+ kv_i = kv_buffer - - # Rearrange query, key, value to (b, s, h, d) - q_i = einops.rearrange(q_i, 's b h d -> b s h d') +- +- # Rearrange query, key, value to (b, s, h, d) +- q_i = einops.rearrange(q_i, 's b h d -> b s h d') - k_i = einops.rearrange(k_i, 's b h d -> b s h d') - v_i = einops.rearrange(v_i, 's b h d -> b s h d') -+ kv_i = einops.rearrange(kv_i, 's b h d -> b s h d') - dout_i = einops.rearrange(dout_i, 's b h d -> b s h d') - - # Backward pass +- dout_i = einops.rearrange(dout_i, 's b h d -> b s h d') +- +- # Backward pass - dq_i, _dk_i, _dv_i, _ = eager_attn_bwd( - q_i, k_i, v_i, attn_bias, None, ctx.scale, ctx.dropout, outs[i], probs[i], dout_i -+ dq_i, _dkv_i, _ = eager_attn_bwd( -+ q_i, kv_i, attn_bias, None, ctx.scale, dim_short, ctx.dropout, outs[i], probs[i], dout_i - ) +- ) ++ comm.all_gather(k_buffer, k) ++ comm.wait() ++ ++ zz_indices = indices.transpose(1, 2) ++ zz_masks = masks.transpose(1, 2) ++ ++ k_i = k_buffer ++ ++ dq_list = [] ++ dk_list = [] ++ ++ s_, b_, h_, d_ = q.shape ++ q = einops.rearrange(q, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = k_i.shape ++ k_i = einops.rearrange(k_i, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = dout.shape ++ dout = einops.rearrange(dout, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ s_, b_, h_, d_ = out.shape ++ out = einops.rearrange(out, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ b_, s_, h_ = lse.shape ++ lse = lse.flatten().view(b_, s_, h_) ++ zz_indices_i = zz_indices ++ b_, s_, g_, topk_ = zz_indices_i.shape ++ zz_indices_i = zz_indices_i.flatten().view(b_, s_, g_, topk_) ++ zz_masks_i = zz_masks ++ b_, s_, g_, skv_ = zz_masks_i.shape ++ zz_masks_i = zz_masks_i.flatten().view(b_, s_, g_, skv_) ++ ++ heads_kv_stride = 16 ++ for i in range(0, heads, heads_kv_stride): ++ q_slice = slice(i, min(i + heads_kv_stride, heads)) ++ q_i = q[:, :, q_slice, :].contiguous() ++ dout_i = dout[:, :, q_slice, :].contiguous() ++ out_i = out[:, :, q_slice, :].contiguous() ++ lse_i = lse[:, :, q_slice].contiguous() ++ ++ # TODO: needs casual = True, may not be compatible with zz ++ dq_i, _dk_i = sparse_mla_bwd(q_i, k_i, out_i, dout_i, zz_indices_i, zz_masks_i, lse_i, dim_v, sm_scale = softmax_scale) - # Rearrange gradients to (s, b, h, d) +- # Rearrange gradients to (s, b, h, d) dq_i = einops.rearrange(dq_i, 'b s h d -> s b h d') -- _dk_i = einops.rearrange(_dk_i, 'b s h d -> s b h d') + _dk_i = einops.rearrange(_dk_i, 'b s h d -> s b h d') - _dv_i = einops.rearrange(_dv_i, 'b s h d -> s b h d') -+ _dkv_i = einops.rearrange(_dkv_i, 'b s h d -> s b h d') + if pg is None: -- dk_i = _dk_i + dk_i = _dk_i - dv_i = _dv_i -+ dkv_i = _dkv_i else: - # Reduce-scatter gradients if CP > 1 -- dk_i = torch.zeros( -- (k_i.shape[1] // cp_size, k_i.shape[0], k_i.shape[2], k_i.shape[3]), -- device=k_i.device, -- dtype=k_i.dtype, -- ) +- # Reduce-scatter gradients if CP > 1 + dk_i = torch.zeros( + (k_i.shape[1] // cp_size, k_i.shape[0], k_i.shape[2], k_i.shape[3]), + device=k_i.device, + dtype=k_i.dtype, + ) - dv_i = torch.zeros( - (v_i.shape[1] // cp_size, v_i.shape[0], v_i.shape[2], v_i.shape[3]), - device=v_i.device, - dtype=v_i.dtype, -+ dkv_i = torch.zeros( -+ (kv_i.shape[1] // cp_size, kv_i.shape[0], kv_i.shape[2], kv_i.shape[3]), -+ device=kv_i.device, -+ dtype=kv_i.dtype, - ) -- torch.distributed.reduce_scatter_tensor(dk_i, _dk_i, group=pg) +- ) + torch.distributed.reduce_scatter_tensor(dk_i, _dk_i, group=pg) - torch.distributed.reduce_scatter_tensor(dv_i, _dv_i, group=pg) -+ torch.distributed.reduce_scatter_tensor(dkv_i, _dkv_i, group=pg) - # Collect gradients - dq.append(dq_i) +- # Collect gradients +- dq.append(dq_i) - dk.append(dk_i) - dv.append(dv_i) -+ dkv.append(dkv_i) ++ dq_list.append(dq_i) ++ dk_list.append(dk_i) # Concatenate gradients and return - dq = torch.cat(dq, dim=2) +- dq = torch.cat(dq, dim=2) - dk = torch.cat(dk, dim=2) - dv = torch.cat(dv, dim=2) - return dq, dk, dv, None, None, None, None -+ dkv = torch.cat(dkv, dim=2) -+ return dq, dkv, dkv[:,:,:,:dim_short].detach().contiguous(), None, None, None, None, None ++ dq = torch.cat(dq_list, dim=2) ++ dk = sum(dk_list) ++ ++ return dq, dk, None, None, None, None, None, None diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py -index fc994490b..7bc9a485e 100644 +index fc994490b..b23d2e9a8 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -6,6 +6,7 @@ from dataclasses import dataclass @@ -358,14 +487,14 @@ index fc994490b..7bc9a485e 100644 + float_mask = torch.zeros_like(causal_mask, dtype=torch.float32).masked_fill( + causal_mask, float('-inf') + ) ++ ++ index_mask = torch.full( ++ (b, sq_local, sk_global), float("-inf"), dtype=torch.float32, device=causal_mask.device ++ ).scatter_(-1, topk_indices, 0) - # Sum attention scores across heads. - # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] - attention_scores = attention_scores.sum(dim=1) -+ index_mask = torch.full( -+ (b, sq_local, sk_global), float("-inf"), dtype=torch.float32, device=causal_mask.device -+ ).scatter_(-1, topk_indices, 0) -+ + float_mask = float_mask.view(1, 1, sq_local, sk_global) + float_mask = index_mask.view(b, 1, sq_local, sk_global) + float_mask if sparse_loss else float_mask + @@ -506,7 +635,7 @@ index fc994490b..7bc9a485e 100644 # [batch, seqlen, index_topk] topk_indices = index_scores.topk(topk_k, dim=-1)[1] -@@ -687,6 +780,57 @@ def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): +@@ -687,6 +780,48 @@ def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): output = output.reshape(sq, b, np * hnv) return output @@ -537,34 +666,56 @@ index fc994490b..7bc9a485e 100644 + + return causal_mask + -+def unfused_dsa_fn_with_cp(query, key, value, topk_indices, softmax_scale): ++def unfused_dsa_fn_with_cp(query, key, dim_v, topk_indices, softmax_scale): + pg = parallel_state.get_context_parallel_group() -+ cp_size = parallel_state.get_context_parallel_world_size() -+ cp_rank = parallel_state.get_context_parallel_rank() -+ + sq, b, np, hn = query.size() + skv = key.size(0) -+ hnv = value.size(3) -+ -+ skv_global = skv * cp_size -+ -+ sparse_mask = torch.ones((b, sq, skv_global), dtype=torch.bool, device=query.device) -+ sparse_mask.scatter_(-1, topk_indices, False) + -+ causal_mask = get_causal_mask(sq, skv, query.device) -+ -+ combined_mask = sparse_mask | causal_mask.unsqueeze(0) -+ -+ attention_mask_for_cp = combined_mask.unsqueeze(1) # [b, 1, sq, skv_global] ++ topk = topk_indices.shape[-1] ++ topk_indices = topk_indices.unsqueeze(1) ++ topk_indices = topk_indices.expand(-1, key.shape[2], -1, -1).contiguous().to(torch.int32) ++ causal_masks = get_causal_mask(sq, skv, query.device) ++ causal_masks = causal_masks[None, None, :, :] ++ causal_masks = causal_masks.expand(b, key.shape[2], -1, -1).contiguous() + output = AttentionFuncionWithContextParallel.apply( -+ query, key, value, attention_mask_for_cp, 0.0, softmax_scale, pg, True ++ query, key, dim_v, topk_indices, causal_masks, 0.0, softmax_scale, pg + ) -+ return output.reshape(sq, b, np * hnv) -+ ++ return output.reshape(sq, b, np * dim_v) class DSAttention(MegatronModule): """ -@@ -768,18 +912,17 @@ class DSAttention(MegatronModule): +@@ -729,7 +864,6 @@ class DSAttention(MegatronModule): + self, + query: torch.Tensor, + key: torch.Tensor, +- value: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + attention_mask: torch.Tensor, +@@ -743,7 +877,6 @@ class DSAttention(MegatronModule): + Args: + query: Query tensor [sq, b, np, hn]. + key: Key tensor [skv, b, np, hn]. +- value: Value tensor [skv, b, np, hnv]. + x: Original hidden states [sq, b, hidden_size]. + qr: Low-rank query representation [sq, b, q_lora_rank]. + attention_mask: Attention mask tensor [b, 1, sq, sk]. +@@ -754,9 +887,11 @@ class DSAttention(MegatronModule): + Returns: + output: Output tensor [sq, b, hidden_size] + """ +- sq, b, np, hn = query.size() +- skv = key.size(0) +- hnv = value.size(3) ++ dim_v = self.config.kv_lora_rank ++ # torch.Size([128, 1, 64, 576]) ++ sq, b, nheads, dim = query.size() ++ # torch.Size([128, 1, 1, 576]) ++ skv, _, kv_groups, _ = key.shape + + # Detach x and qr to prevent gradients of indexer from flowing back to the main model. + x = x.detach() +@@ -768,18 +903,17 @@ class DSAttention(MegatronModule): # Generate upper triangular mask with -inf above diagonal, 0 elsewhere # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal) # float_mask [sq, skv] @@ -591,17 +742,42 @@ index fc994490b..7bc9a485e 100644 # =================================== # Get index scores and top-k indices -@@ -791,7 +934,7 @@ class DSAttention(MegatronModule): +@@ -791,32 +925,6 @@ class DSAttention(MegatronModule): # =================================== # Run sparse attention kernel # =================================== - output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) -+ output = unfused_dsa_fn_with_cp(query, key, value, topk_indices, self.softmax_scale) +- +- # =================================== +- # Attach indexer loss +- # =================================== +- if self.training and torch.is_grad_enabled(): +- # Compute KL divergence loss between indexer scores and true attention scores +- indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) +- indexer_loss = compute_dsa_indexer_loss( +- index_scores, +- topk_indices, +- query.detach(), +- key.detach(), +- self.softmax_scale, +- indexer_loss_coeff, +- getattr(self.config, "dsa_indexer_use_sparse_loss", False), +- self.indexer.pg_collection, +- ) +- # Save indexer loss for logging +- if indexer_loss_coeff > 0: +- DSAIndexerLossLoggingHelper.save_loss_to_tracker( +- loss=indexer_loss, +- layer_number=self.layer_number, +- num_layers=self.config.num_layers, +- ) +- # Attach loss to output +- output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) ++ output = unfused_dsa_fn_with_cp(query, key, dim_v, topk_indices, self.softmax_scale) - # =================================== - # Attach indexer loss + return output diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py -index 3953d933b..84301ed54 100644 +index 3953d933b..7d030ad02 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -6,6 +6,7 @@ from dataclasses import dataclass @@ -612,6 +788,692 @@ index 3953d933b..84301ed54 100644 try: from einops import rearrange +@@ -167,6 +168,7 @@ class MultiLatentAttention(Attention): + ) + + # Output. ++ # SP_Reduce scatter + TP_Row_par + self.linear_proj = build_module( + submodules.linear_proj, + self.query_projection_size, +@@ -311,7 +313,6 @@ class MultiLatentAttention(Attention): + core_attn_out = self.core_attention( + query, + key, +- value, + x=hidden_states, + qr=q_compressed, + attention_mask=attention_mask, +@@ -370,6 +371,19 @@ class MultiLatentAttention(Attention): + self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) + self.qkv_up_checkpoint = None + ++ s_, b_ = core_attn_out.size(0), core_attn_out.size(1) ++ core_attn_out = core_attn_out.view( ++ s_, b_, ++ self.num_attention_heads_per_partition, ++ self.config.kv_lora_rank ++ ) ++ ++ # einsum: "sbhk,hdk->sbhd" ++ core_attn_out = torch.einsum("sbhk,hdk->sbhd", core_attn_out, self.up_v_weight_) ++ core_attn_out = core_attn_out.contiguous() ++ core_attn_out = core_attn_out.view(s_, b_, -1) ++ core_attn_out = core_attn_out.contiguous() ++ + # ================= + # Output. [sq, b, h] + # ================= +@@ -384,7 +398,6 @@ class MultiLatentAttention(Attention): + + return output, bias + +- + class MLASelfAttention(MultiLatentAttention): + """MLA Self-attention layer class + +@@ -753,7 +766,6 @@ class MLASelfAttention(MultiLatentAttention): + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + +- # todo add assert about fusions and caching + if self.config.apply_rope_fusion: + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() +@@ -844,6 +856,98 @@ class MLASelfAttention(MultiLatentAttention): + value = value.contiguous() + + return query, key, value ++ ++ def mla_absorb(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): ++ if self.config.q_lora_rank is not None: ++ # q_compressed: [num_tokens, q_lora_rank] ++ # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] ++ q, _ = self.linear_q_up_proj(q_compressed) ++ else: ++ # q_compressed: [num_tokens, hidden_size] ++ # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] ++ q, _ = self.linear_q_proj(q_compressed) ++ ++ # q: [num_tokens, n, q_head_dim] ++ q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) ++ ++ # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] ++ k_pos_emb = torch.unsqueeze(k_pos_emb, -2) ++ ++ if self.config.apply_rope_fusion: ++ raise NotImplementedError( ++ "RoPE fusion is not yet supported with absorption training. " ++ "Please set apply_rope_fusion=False." ++ ) ++ else: ++ q_len = q.size()[0] ++ if inference_context is not None: ++ # add offset to the sequence start for inference ++ sequence_start = inference_context.sequence_len_offset ++ sequence_end = sequence_start + q_len ++ rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end] ++ elif packed_seq_params is None or self.config.context_parallel_size == 1: ++ rotary_pos_emb = rotary_pos_emb[0:q_len] ++ ++ # q_no_pe: [num_tokens, n, qk_head_dim] ++ # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] ++ q_no_pe, q_pos_emb = torch.split( ++ q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 ++ ) ++ ++ # q_no_pe: [num_tokens, n, qk_head_dim] ++ # up_k_weight: [n, qk_head_dim, kv_lora_rank] ++ # q_absorbed: [num_tokens, n, kv_lora_rank] ++ q_absorbed = torch.einsum("...hd,hdk->...hk", q_no_pe, self.up_k_weight_) ++ ++ # TODO: Does it match ZZ? SP does not need but CP needs ++ if self.config.sequence_parallel: ++ kv_compressed = gather_from_sequence_parallel_region(kv_compressed) ++ ++ # kv_compressed: [num_tokens, kv_lora_rank] ++ if kv_compressed.ndim == 3: # [s, b, kv_lora_rank] ++ k_content = kv_compressed.unsqueeze(2).expand( ++ -1, -1, 1, -1 ++ ) ++ else: # [t, kv_lora_rank] for packed sequence ++ k_content = kv_compressed.unsqueeze(1).expand( ++ -1, 1, -1 ++ ) ++ ++ # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] ++ q_pos_emb = apply_rotary_pos_emb( ++ q_pos_emb, ++ rotary_pos_emb, ++ config=self.config, ++ cu_seqlens=cu_seqlens_q, ++ mscale=mscale, ++ cp_group=self.pg_collection.cp, ++ ) ++ # k_pos_emb: [num_tokens, 1, qk_pos_emb_head_dim] ++ k_pos_emb = apply_rotary_pos_emb( ++ k_pos_emb, ++ rotary_pos_emb, ++ config=self.config, ++ cu_seqlens=cu_seqlens_kv, ++ mscale=mscale, ++ cp_group=self.pg_collection.cp, ++ ) ++ ++ # query: [num_tokens, n, kv_lora_rank + qk_pos_emb_head_dim] ++ query = torch.cat([q_absorbed, q_pos_emb], dim=-1) ++ ++ # key: [num_tokens, n, kv_lora_rank + qk_pos_emb_head_dim] ++ if k_pos_emb.ndim == 4: ++ k_pos_emb = k_pos_emb.expand(-1, -1, 1, -1) ++ else: ++ assert k_pos_emb.ndim == 3 ++ k_pos_emb = k_pos_emb.expand(-1, 1, -1) ++ ++ key = torch.cat([k_content, k_pos_emb], dim=-1) ++ ++ query = query.contiguous() ++ key = key.contiguous() ++ ++ return query, key + + if self.recompute_up_proj: + quantization = self.config.fp8 or self.config.fp4 +@@ -860,9 +964,10 @@ class MLASelfAttention(MultiLatentAttention): + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + else: +- query, key, value = qkv_up_proj_and_rope_apply( ++ query, key = mla_absorb( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) ++ value = None + + if return_compressed_tensors: + return query, key, value, q_compressed, kv_compressed +@@ -1104,5 +1209,27 @@ class MLASelfAttention(MultiLatentAttention): + * (self.config.qk_head_dim + self.config.v_head_dim), + -1, + ) +- + return weight_kv_updated ++ ++ @property ++ def up_k_weight_(self): ++ # linear_kv_up_proj.weight: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank] ++ weight = self.linear_kv_up_proj.weight ++ weight_reshaped = weight.view( ++ self.num_attention_heads_per_partition, ++ self.config.qk_head_dim + self.config.v_head_dim, ++ self.config.kv_lora_rank, ++ ) ++ # [num_heads_per_partition, qk_head_dim, kv_lora_rank] ++ return weight_reshaped[:, :self.config.qk_head_dim, :] ++ ++ @property ++ def up_v_weight_(self): ++ weight = self.linear_kv_up_proj.weight ++ weight_reshaped = weight.view( ++ self.num_attention_heads_per_partition, ++ self.config.qk_head_dim + self.config.v_head_dim, ++ self.config.kv_lora_rank, ++ ) ++ # [num_heads_per_partition, v_head_dim, kv_lora_rank] ++ return weight_reshaped[:, self.config.qk_head_dim:, :] +diff --git a/megatron/core/transformer/tilelang_kernel/__init__.py b/megatron/core/transformer/tilelang_kernel/__init__.py +new file mode 100644 +index 000000000..d8f2425f0 +--- /dev/null ++++ b/megatron/core/transformer/tilelang_kernel/__init__.py +@@ -0,0 +1,10 @@ ++# Code is adopted from tilelang/examples/deepseek_v32 ++# transformer/tilelang_kernel/__init__.py ++ ++from .sparse_mla_fwd import sparse_mla_fwd_interface ++from .sparse_mla_bwd import sparse_mla_bwd ++ ++__all__ = [ ++ "sparse_mla_fwd_interface", ++ "sparse_mla_bwd", ++] +diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py +new file mode 100644 +index 000000000..b8ea416dd +--- /dev/null ++++ b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py +@@ -0,0 +1,274 @@ ++# ruff: noqa ++import tilelang ++from tilelang import language as T ++import torch ++ ++ ++@tilelang.jit(out_idx=[-1]) ++def preprocess( ++ B, ++ S, ++ H, ++ D, ++ block_ND=32, ++ num_stages=5, ++ dtype=T.bfloat16, ++ accum_dtype=T.float32, ++): ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ shape = [B, S, H, D] ++ ++ @T.prim_func ++ def preprocess_kernel( ++ O: T.Tensor(shape, dtype), ++ dO: T.Tensor(shape, dtype), ++ Delta: T.Tensor([B, S, H], accum_dtype), ++ ): ++ with T.Kernel(H, T.ceildiv(S, block_ND), B) as (bx, by, bz): ++ o = T.alloc_fragment([block_ND, block_ND], accum_dtype) ++ do = T.alloc_fragment([block_ND, block_ND], accum_dtype) ++ delta = T.alloc_fragment([block_ND], accum_dtype) ++ acc = T.alloc_fragment([block_ND, block_ND], accum_dtype) ++ T.clear(acc) ++ for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages): ++ T.copy(O[bz, by * block_ND : (by + 1) * block_ND, bx, k * block_ND : (k + 1) * block_ND], o) ++ T.copy(dO[bz, by * block_ND : (by + 1) * block_ND, bx, k * block_ND : (k + 1) * block_ND], do) ++ for i, j in T.Parallel(block_ND, block_ND): ++ acc[i, j] += o[i, j] * do[i, j] ++ T.reduce_sum(acc, delta, 1) ++ T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx]) ++ ++ return preprocess_kernel ++ ++ ++@tilelang.jit(out_idx=[-1]) ++def postprocess( ++ B, ++ S_kv, ++ D, ++ D_tail, ++ kv_group=1, ++ block_N=64, ++ threads=256, ++ dtype=T.bfloat16, ++ accum_dtype=T.float32, ++): ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ dkv_shape = [B, S_kv, kv_group, D + D_tail] ++ ++ @T.prim_func ++ def postprocess_kernel( ++ dKV: T.Tensor(dkv_shape, accum_dtype), ++ dKV_out: T.Tensor(dkv_shape, dtype), ++ ): ++ with T.Kernel(T.ceildiv(S_kv, block_N), kv_group, B, threads=threads) as (bx, by, bz): ++ T.copy( ++ dKV[bz, bx * block_N : (bx + 1) * block_N, by, :], ++ dKV_out[bz, bx * block_N : (bx + 1) * block_N, by, :], ++ ) ++ ++ return postprocess_kernel ++ ++ ++@tilelang.jit( ++ out_idx=[-2], ++ pass_configs={ ++ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, ++ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, ++ tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True, ++ }, ++) ++def bwd( ++ B, ++ S, ++ S_kv, ++ H, ++ D, ++ D_tail, ++ topk, ++ kv_group=1, ++ sm_scale=None, ++ is_causal=True, ++ block_size=32, ++ num_stages=0, ++ threads=128, ++ indices_dtype=T.int32, ++ dtype=T.bfloat16, ++ accum_dtype=T.float32, ++ masks_dtype=T.bool, ++): ++ assert is_causal == True, "non-casual is not supported now" ++ assert topk % block_size == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded" ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ assert indices_dtype == T.int32 ++ ++ if sm_scale is None: ++ sm_scale = (D + D_tail) ** (-0.5) ++ sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e) ++ ++ H_kv = H // kv_group ++ q_shape = [B, S, H, D + D_tail] ++ k_shape = [B, S_kv, kv_group, D + D_tail] ++ o_shape = [B, S, H, D] ++ indices_shape = [B, S, kv_group, topk] ++ delta_shape = [B, S, H] ++ lse_shape = [B, S, H] ++ masks_shape = [B, S, kv_group, S_kv] ++ assert indices_dtype == T.int32 ++ assert dtype == T.bfloat16 ++ assert accum_dtype == T.float32 ++ ++ H = H_kv ++ padded_H = max(tilelang.math.next_power_of_2(H_kv), 16) ++ block_H = min(64, padded_H) ++ assert padded_H % block_H == 0 ++ NH = padded_H // block_H ++ BS = block_size ++ NS = tilelang.cdiv(topk, block_size) ++ ++ split_store = 2 ++ ++ @T.prim_func ++ def sparse_mla_bwd_kernel( ++ Q: T.Tensor(q_shape, dtype), ++ KV: T.Tensor(k_shape, dtype), ++ dO: T.Tensor(o_shape, dtype), ++ Indices: T.Tensor(indices_shape, indices_dtype), ++ Masks: T.Tensor(masks_shape, masks_dtype), ++ Lse: T.Tensor(lse_shape, accum_dtype), ++ Delta: T.Tensor(delta_shape, accum_dtype), ++ dQ: T.Tensor(q_shape, dtype), ++ dKV: T.Tensor(k_shape, accum_dtype), ++ ): ++ with T.Kernel(S, B, kv_group * NH, threads=threads) as (s_i, by, bz): ++ Q_shared = T.alloc_shared([block_H, D], dtype) ++ Q_tail_shared = T.alloc_shared([block_H, D_tail], dtype) ++ KV_shared = T.alloc_shared([BS, D], dtype) ++ KV_tail_shared = T.alloc_shared([BS, D_tail], dtype) ++ dO_shared = T.alloc_shared([block_H, D], dtype) ++ mask = T.alloc_fragment([BS], "bool") ++ ++ P_shared_cast = T.alloc_shared([block_H, BS], dtype) ++ dP_shared_cast = T.alloc_shared([block_H, BS], dtype) ++ dQ_shared = T.alloc_shared([block_H, D], dtype) ++ dQ_tail_shared = T.alloc_shared([block_H, D_tail], dtype) ++ ++ acc_p = T.alloc_fragment([block_H, BS], accum_dtype) ++ acc_dp = T.alloc_fragment([block_H, BS], accum_dtype) ++ acc_dq = T.alloc_fragment([block_H, D], accum_dtype) ++ acc_dq_tail = T.alloc_fragment([block_H, D_tail], accum_dtype) ++ acc_dkv = T.alloc_fragment([BS, D], accum_dtype) ++ acc_dkv_tail = T.alloc_fragment([BS, D_tail], accum_dtype) ++ acc_dkv_shared = T.alloc_shared([BS // split_store, D], accum_dtype) ++ acc_dkv_tail_shared = T.alloc_shared([BS // split_store, D_tail], accum_dtype) ++ ++ T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, :D], Q_shared) ++ T.copy(Q[by, s_i, bz * block_H : (bz + 1) * block_H, D:], Q_tail_shared) ++ T.copy(dO[by, s_i, bz * block_H : (bz + 1) * block_H, :D], dO_shared) ++ ++ T.clear(acc_dq) ++ T.clear(acc_dq_tail) ++ ++ # Process each block of indices ++ for i_i in T.Pipelined(NS, num_stages=num_stages): ++ # Compute attention scores ++ for bi_i in T.Parallel(BS): ++ mask[bi_i] = Masks[by, s_i, bz // NH, Indices[by, s_i, bz // NH, i_i * BS + bi_i]] ++ ++ for h_i, bi_i in T.Parallel(block_H, BS): ++ acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], -T.infinity(acc_p.dtype), 0) ++ ++ # Load KV, V for this block of indices ++ for bi_i, d_i in T.Parallel(BS, D): ++ KV_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, d_i] ++ ++ T.gemm(Q_shared, KV_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ for bi_i, d_i in T.Parallel(BS, D_tail): ++ KV_tail_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, D + d_i] ++ T.gemm(Q_tail_shared, KV_tail_shared[:, :D_tail], acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ for h_i, bi_i in T.Parallel(block_H, BS): ++ acc_p[h_i, bi_i] = T.exp2(acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2 - Lse[by, s_i, bz * block_H + h_i]) ++ ++ T.copy(acc_p, P_shared_cast) ++ ++ T.gemm(dO_shared, KV_shared, acc_dp, transpose_B=True, policy=T.GemmWarpPolicy.FullCol, clear_accum=True) ++ ++ for h_i, bi_i in T.Parallel(block_H, BS): ++ acc_dp[h_i, bi_i] = acc_p[h_i, bi_i] * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * block_H + h_i]) * sm_scale ++ ++ T.copy(acc_dp, dP_shared_cast) ++ T.gemm(dP_shared_cast, KV_shared, acc_dq, policy=T.GemmWarpPolicy.FullCol) ++ T.gemm(dP_shared_cast, KV_tail_shared, acc_dq_tail, policy=T.GemmWarpPolicy.FullCol) ++ ++ T.gemm(dP_shared_cast, Q_shared, acc_dkv, transpose_A=True, policy=T.GemmWarpPolicy.FullCol, clear_accum=True) ++ T.gemm(P_shared_cast, dO_shared, acc_dkv, transpose_A=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ T.clear(acc_dkv_tail) ++ T.gemm(dP_shared_cast, Q_tail_shared, acc_dkv_tail, transpose_A=True, policy=T.GemmWarpPolicy.FullCol) ++ ++ for s in range(split_store): ++ for bi_i, d_i in T.Parallel(BS, D): ++ if bi_i < BS // split_store: ++ acc_dkv_shared[bi_i, d_i] = acc_dkv[bi_i + s * (BS // split_store), d_i] ++ ++ for bi_i, d_i in T.Parallel(BS, D_tail): ++ if bi_i < BS // split_store: ++ acc_dkv_tail_shared[bi_i, d_i] = acc_dkv_tail[bi_i + s * (BS // split_store), d_i] ++ ++ for bi_i, d_i in T.Parallel(BS // split_store, D // 4): ++ T.atomic_addx4( ++ dKV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], bz // NH, d_i * 4], ++ acc_dkv_shared[bi_i, d_i * 4], ++ ) ++ ++ # Atomically update dKV, dKV_tail tensors ++ for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): ++ T.atomic_addx4( ++ dKV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], bz // NH, D + d_i * 4], ++ acc_dkv_tail_shared[bi_i, d_i * 4], ++ ) ++ ++ # Store the accumulated dQ ++ T.copy(acc_dq, dQ_shared) ++ T.copy(acc_dq_tail[:, :D_tail], dQ_tail_shared) ++ ++ T.copy(dQ_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, :D]) ++ T.copy(dQ_tail_shared, dQ[by, s_i, bz * block_H : (bz + 1) * block_H, D:]) ++ ++ return sparse_mla_bwd_kernel ++ ++ ++def sparse_mla_bwd(q, kv, o, do, indices, masks, lse, dim_v, sm_scale=None, is_casual=True, return_kernel=False, delta=None): ++ assert q.is_contiguous() ++ assert kv.is_contiguous() ++ assert indices.is_contiguous() ++ assert lse.is_contiguous() ++ B, S, H, dim_plus_tail_dim = q.shape ++ _, S_kv, kv_group, _ = kv.shape ++ assert kv.shape[-1] == dim_plus_tail_dim ++ assert kv.shape[0] == B ++ # dim should be assigned ++ D = dim_v ++ ++ D_tail = dim_plus_tail_dim - D ++ topk = indices.shape[-1] ++ assert indices.shape == (B, S, kv_group, topk) ++ assert lse.shape == (B, S, H) ++ ++ # Get kernels ++ preprocess_kernel = preprocess(B, S, H, D) ++ bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual) ++ postprocess_kernel = postprocess(B, S_kv, D, D_tail, kv_group) ++ ++ if delta is None: ++ delta = preprocess_kernel(o, do) ++ dkv = torch.zeros_like(kv, dtype=torch.float32) ++ dq = bwd_kernel(q, kv, do, indices, masks, lse, delta, dkv) ++ dkv = postprocess_kernel(dkv) ++ ++ return dq, dkv +\ No newline at end of file +diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py +new file mode 100644 +index 000000000..d338a2fa6 +--- /dev/null ++++ b/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py +@@ -0,0 +1,190 @@ ++# ruff: noqa ++import torch ++import tilelang ++from tilelang import language as T ++ ++ ++@tilelang.jit( ++ out_idx=[-2, -1], ++ pass_configs={ ++ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, ++ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, ++ }, ++) ++def sparse_mla_fwd( ++ heads, ++ dim, ++ tail_dim, ++ topk, ++ kv_group=1, ++ sm_scale=None, ++ is_causal=True, ++ CP0=True, ++ block_I=64, ++ num_stages=2, ++ threads=256, ++): ++ assert dim == tilelang.math.next_power_of_2(dim), f"haven't check padding correctness yet, dim={dim}" ++ assert tail_dim == tilelang.math.next_power_of_2(tail_dim), f"haven't check padding correctness yet, dim={tail_dim}" ++ assert is_causal == True, "non-casual is not supported" ++ assert topk % block_I == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded" ++ if sm_scale is None: ++ sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) ++ else: ++ sm_scale = sm_scale * 1.44269504 # log2(e) ++ ++ batch = T.dynamic("batch") ++ seq_len = T.dynamic("seq_len") ++ seq_len_kv = T.dynamic("seq_len_kv") ++ ++ head_kv = heads // kv_group ++ q_shape = [batch, seq_len, heads, dim + tail_dim] ++ kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] ++ o_shape = [batch, seq_len, heads, dim] ++ indices_shape = [batch, seq_len, kv_group, topk] ++ lse_shape = [batch, seq_len, heads] ++ masks_shape = [batch, seq_len, kv_group, seq_len_kv] ++ ++ masks_dtype = T.bool ++ indices_dtype = T.int32 ++ dtype = T.bfloat16 ++ accum_dtype = T.float32 ++ ++ G = kv_group ++ H = head_kv ++ padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) ++ if padded_H != H: ++ assert kv_group == 1, ( ++ "here we solve the H padding automatically, other wise you should handle Q copy and Output copy with your mask (when kv_group == 1, use g_i * padded_H:(g_i+1) * padded_H would be handled automatically)" ++ ) ++ BI = block_I ++ NI = tilelang.cdiv(topk, block_I) ++ D = dim ++ D_tail = tail_dim ++ ++ if head_kv > 64: ++ assert head_kv % 64 == 0, "head_kv should be a multiple of 64" ++ REPLICATE_H = head_kv // 64 ++ else: ++ REPLICATE_H = 1 ++ ++ H_per_block = padded_H if REPLICATE_H == 1 else 64 ++ ++ @T.prim_func ++ def main( ++ Q: T.Tensor(q_shape, dtype), # type: ignore ++ KV: T.Tensor(kv_shape, dtype), # type: ignore ++ Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore ++ Masks: T.Tensor(masks_shape, masks_dtype), # type: ignore ++ Output: T.Tensor(o_shape, dtype), # type: ignore ++ Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore ++ ): ++ with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as ( ++ bx, ++ by, ++ bz, ++ ): ++ Q_shared = T.alloc_shared([H_per_block, D], dtype) ++ Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) ++ KV_shared = T.alloc_shared([BI, D], dtype) ++ K_tail_shared = T.alloc_shared([BI, D_tail], dtype) ++ O_shared = T.alloc_shared([H_per_block, D], dtype) ++ Lse_shared = T.alloc_shared([H_per_block], accum_dtype) ++ mask = T.alloc_fragment([BI], "bool") ++ ++ acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) ++ acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) ++ S_shared = T.alloc_shared([H_per_block, BI], dtype) ++ sumexp = T.alloc_fragment([H_per_block], accum_dtype) ++ sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) ++ alpha = T.alloc_fragment([H_per_block], accum_dtype) ++ m_i = T.alloc_fragment([H_per_block], accum_dtype) ++ m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) ++ ++ T.fill(acc_o, 0) ++ T.fill(sumexp, 0) ++ T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan ++ ++ b_i, g_i = by, bz ++ s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) ++ q_i = s_i ++ ++ H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) ++ H1 = H0 + H_per_block ++ ++ T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) ++ T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) ++ ++ for i_i in T.Pipelined(NI, num_stages=num_stages): ++ for bi_i in T.Parallel(BI): ++ mask[bi_i] = Masks[b_i, s_i, g_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i]] ++ for bi_i, d_i in T.Parallel(BI, D): ++ KV_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] ++ for bi_i, d_i in T.Parallel(BI, D_tail): ++ K_tail_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i] ++ for h_i, bi_i in T.Parallel(H_per_block, BI): ++ acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], -T.infinity(acc_s.dtype), 0) ++ T.gemm( ++ Q_shared, ++ KV_shared, ++ acc_s, ++ transpose_B=True, ++ policy=T.GemmWarpPolicy.FullRow, ++ ) ++ T.gemm( ++ Q_tail_shared, ++ K_tail_shared, ++ acc_s, ++ transpose_B=True, ++ policy=T.GemmWarpPolicy.FullRow, ++ ) ++ T.copy(m_i, m_i_prev) ++ T.reduce_max(acc_s, m_i, dim=1, clear=False) ++ for h_i in T.Parallel(H_per_block): ++ m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) ++ for h_i in T.Parallel(H_per_block): ++ alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) ++ for h_i, bi_i in T.Parallel(H_per_block, BI): ++ acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale) ++ T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? ++ for h_i in T.Parallel(H_per_block): ++ sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] ++ for h_i, d_i in T.Parallel(H_per_block, D): ++ acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] ++ ++ T.copy(acc_s, S_shared) ++ T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) ++ ++ # Rescale ++ for h_i, d_i in T.Parallel(H_per_block, D): ++ acc_o[h_i, d_i] /= sumexp[h_i] ++ for h_i in T.Parallel(H_per_block): ++ sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale ++ ++ T.copy(acc_o, O_shared) ++ T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) ++ T.copy(sumexp, Lse_shared) ++ T.copy(sumexp, Lse[b_i, s_i, H0:H1]) ++ ++ return main ++ ++ ++def sparse_mla_fwd_interface(q, kv, indices, masks, d_v, sm_scale=None, return_p_sum: bool = False, block_I=64, num_stages=2, threads=256): ++ is_casual = True ++ assert return_p_sum == False, "This kernel file is for fwd only" ++ assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous() ++ batch, seq_len, heads, dim_plus_tail_dim = q.shape ++ _, seq_len_kv, kv_group, _ = kv.shape ++ ++ assert kv.shape[-1] == dim_plus_tail_dim ++ tail_dim = dim_plus_tail_dim - d_v ++ assert kv.shape[0] == batch ++ _, _, _, topk = indices.shape ++ assert indices.shape == (batch, seq_len, kv_group, topk) ++ assert masks.shape == (batch, seq_len, kv_group, seq_len_kv) ++ ++ kernel = sparse_mla_fwd( ++ heads, d_v, tail_dim, topk, kv_group, sm_scale, is_casual, block_I=block_I, num_stages=num_stages, threads=threads ++ ) ++ out, lse = kernel(q, kv, indices, masks) ++ return out, lse +\ No newline at end of file diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index a3a167549..98391fda6 100644 --- a/megatron/core/transformer/transformer_config.py diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index 87cf24992d7..da5b2b55ad5 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -43,7 +43,9 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio if converted_name.endswith("_scale"): continue if_use_ue8m0_in_moe = True if args.sglang_moe_runner_backend == "deep_gemm" else False - quantize_named_params.extend(_quantize_param(converted_name, param, weight_block_size, if_use_ue8m0_in_moe=if_use_ue8m0_in_moe)) + quantize_named_params.extend( + _quantize_param(converted_name, param, weight_block_size, if_use_ue8m0_in_moe=if_use_ue8m0_in_moe) + ) return quantize_named_params @@ -92,9 +94,11 @@ def _quantize_param(name, weight, weight_block_size, if_use_ue8m0_in_moe=True): FP8_MIN = torch.finfo(torch.float8_e4m3fn).min FP8_MAX = torch.finfo(torch.float8_e4m3fn).max if weight_block_size is not None: - if should_deepgemm_weight_requant_ue8m0 and should_deepgemm_weight_requant_ue8m0( - weight_block_size=weight_block_size - ) and if_use_ue8m0_in_moe: + if ( + should_deepgemm_weight_requant_ue8m0 + and should_deepgemm_weight_requant_ue8m0(weight_block_size=weight_block_size) + and if_use_ue8m0_in_moe + ): qweight, scale = quant_weight_ue8m0(weight, weight_block_size=weight_block_size) scale = transform_scale_ue8m0(scale, mn=qweight.shape[-2]) else: diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index 558a2e06f0c..e958566dcaa 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -202,11 +202,17 @@ def _named_params_and_buffers_global( expert_idx = int(expert_idx) + expert_offset yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param continue - + # TODO: a hacking here, need to be cleaner - duplicated = ['indexer.linear_weights_proj', 'indexer.linear_wk', 'indexer.linear_wq_b', 'linear_q_down_proj', 'linear_kv_down_proj'] + duplicated = [ + "indexer.linear_weights_proj", + "indexer.linear_wk", + "indexer.linear_wq_b", + "linear_q_down_proj", + "linear_kv_down_proj", + ] if any(dup in name for dup in duplicated): - param.parallel_mode = 'duplicated' + param.parallel_mode = "duplicated" layer_idx, rest = match.groups() layer_idx = int(layer_idx) + layer_offset diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index cc42522eeec..67b824aa94e 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -1,6 +1,7 @@ -import sys import os -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../..")) from .deepseekv32 import DeepseekV32Bridge from .glm4 import GLM4Bridge @@ -14,12 +15,15 @@ _original_from_config = AutoBridge.from_config + @classmethod def _patched_from_config(cls, hf_config, **kwargs): - if hasattr(hf_config, 'index_n_heads'): + if hasattr(hf_config, "index_n_heads"): from mbridge.core.bridge import _MODEL_REGISTRY - return _MODEL_REGISTRY['deepseek_v32'](hf_config, **kwargs) - + + return _MODEL_REGISTRY["deepseek_v32"](hf_config, **kwargs) + return _original_from_config(hf_config, **kwargs) + AutoBridge.from_config = _patched_from_config diff --git a/miles_plugins/mbridge/deepseekv32.py b/miles_plugins/mbridge/deepseekv32.py index aae07ee5321..19e417fa06c 100644 --- a/miles_plugins/mbridge/deepseekv32.py +++ b/miles_plugins/mbridge/deepseekv32.py @@ -1,6 +1,7 @@ +from megatron.core.transformer.enums import AttnBackend + from mbridge.core import register_model from mbridge.models import DeepseekV3Bridge -from megatron.core.transformer.enums import AttnBackend @register_model("deepseek_v32") @@ -13,49 +14,44 @@ class DeepseekV32Bridge(DeepseekV3Bridge): "self_attention.core_attention.indexer.linear_weights_proj.weight", } - _ATTENTION_MAPPING = ( - DeepseekV3Bridge._ATTENTION_MAPPING.copy() - ) - + _ATTENTION_MAPPING = DeepseekV3Bridge._ATTENTION_MAPPING.copy() + # Because the indexer needs the norm output, we cannot use the fused transformer engine impl and have to compute it separately. if "self_attention.linear_q_up_proj.layer_norm_weight" in _ATTENTION_MAPPING: del _ATTENTION_MAPPING["self_attention.linear_q_up_proj.layer_norm_weight"] if "self_attention.linear_kv_up_proj.layer_norm_weight" in _ATTENTION_MAPPING: del _ATTENTION_MAPPING["self_attention.linear_kv_up_proj.layer_norm_weight"] - - _ATTENTION_MAPPING.update({ - "self_attention.q_layernorm.weight": [ - "model.layers.{layer_number}.self_attn.q_a_layernorm.weight" - ], - "self_attention.kv_layernorm.weight": [ - "model.layers.{layer_number}.self_attn.kv_a_layernorm.weight" - ], - "self_attention.core_attention.indexer.linear_wq_b.weight": [ - "model.layers.{layer_number}.self_attn.indexer.wq_b.weight" - ], - "self_attention.core_attention.indexer.linear_wk.weight": [ - "model.layers.{layer_number}.self_attn.indexer.wk.weight" - ], - "self_attention.core_attention.indexer.k_norm.weight": [ - "model.layers.{layer_number}.self_attn.indexer.k_norm.weight" - ], - "self_attention.core_attention.indexer.k_norm.bias": [ - "model.layers.{layer_number}.self_attn.indexer.k_norm.bias" - ], - "self_attention.core_attention.indexer.linear_weights_proj.weight": [ - "model.layers.{layer_number}.self_attn.indexer.weights_proj.weight" - ], - }) + + _ATTENTION_MAPPING.update( + { + "self_attention.q_layernorm.weight": ["model.layers.{layer_number}.self_attn.q_a_layernorm.weight"], + "self_attention.kv_layernorm.weight": ["model.layers.{layer_number}.self_attn.kv_a_layernorm.weight"], + "self_attention.core_attention.indexer.linear_wq_b.weight": [ + "model.layers.{layer_number}.self_attn.indexer.wq_b.weight" + ], + "self_attention.core_attention.indexer.linear_wk.weight": [ + "model.layers.{layer_number}.self_attn.indexer.wk.weight" + ], + "self_attention.core_attention.indexer.k_norm.weight": [ + "model.layers.{layer_number}.self_attn.indexer.k_norm.weight" + ], + "self_attention.core_attention.indexer.k_norm.bias": [ + "model.layers.{layer_number}.self_attn.indexer.k_norm.bias" + ], + "self_attention.core_attention.indexer.linear_weights_proj.weight": [ + "model.layers.{layer_number}.self_attn.indexer.weights_proj.weight" + ], + } + ) def _build_config(self): config = super()._build_config() - + config.attention_backend = AttnBackend.auto - + config.experimental_attention_variant = "dsa" - config.dsa_indexer_n_heads = getattr(self.hf_config, 'dsa_indexer_n_heads', 64) - config.dsa_indexer_head_dim = getattr(self.hf_config, 'dsa_indexer_head_dim', 128) - config.dsa_indexer_topk = getattr(self.hf_config, 'dsa_indexer_topk', 2048) - - return config + config.dsa_indexer_n_heads = getattr(self.hf_config, "dsa_indexer_n_heads", 64) + config.dsa_indexer_head_dim = getattr(self.hf_config, "dsa_indexer_head_dim", 128) + config.dsa_indexer_topk = getattr(self.hf_config, "dsa_indexer_topk", 2048) + return config diff --git a/scripts/run_deepseek_v32.py b/scripts/run_deepseek_v32.py index 0cbf9ca05d5..6d76ade8611 100644 --- a/scripts/run_deepseek_v32.py +++ b/scripts/run_deepseek_v32.py @@ -5,7 +5,6 @@ import re from dataclasses import dataclass from typing import Literal -from pathlib import Path import typer import miles.utils.external_utils.command_utils as U @@ -282,7 +281,7 @@ def train(args: ScriptArgs): "--use-fault-tolerance " f"--dump-details /root/shared_data/{args.run_id}/dump_details " "--disable-weights-backuper " - "--model-name deepseekv32 " # for mbridge load + "--model-name deepseekv32 " # for mbridge load "--train-memory-margin-bytes 1073741824 " # "--check-weight-update-equal " "--qkv-format bshd " @@ -313,4 +312,4 @@ def train(args: ScriptArgs): if __name__ == "__main__": - app() \ No newline at end of file + app() From 2c3534b66e616d8bd7eb2cd3f07cf241f841e704 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Mon, 19 Jan 2026 00:21:41 -0800 Subject: [PATCH 26/30] update --- docker/deepseekv32/megatron.patch | 64 +++++++++++++++++-------------- 1 file changed, 35 insertions(+), 29 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index ac7a1be3c1c..c1f80e1e50d 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -1,8 +1,8 @@ diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py -index 89659a1d7..1def27c69 100644 +index 89659a1d7..2c1464fb6 100644 --- a/megatron/core/transformer/dot_product_attention_context_parallel.py +++ b/megatron/core/transformer/dot_product_attention_context_parallel.py -@@ -3,9 +3,12 @@ +@@ -3,107 +3,12 @@ # Some of this code was adopted from https://github.com/zhuzilin/ring-flash-attention/ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. @@ -11,14 +11,15 @@ index 89659a1d7..1def27c69 100644 import torch +import torch.distributed as dist from torch.nn import functional as F -+from .tilelang_kernel import sparse_mla_bwd, sparse_mla_fwd_interface - - try: - import einops -@@ -15,96 +18,6 @@ except ImportError: - HAVE_EINOPS = False - - +- +-try: +- import einops +- +- HAVE_EINOPS = True +-except ImportError: +- HAVE_EINOPS = False +- +- -@torch.no_grad -def eager_attn_fwd(q, k, v, attn_bias, sinks, scale, dropout): - """Forward pass for eager attention""" @@ -108,11 +109,11 @@ index 89659a1d7..1def27c69 100644 - grad_q = einops.rearrange(grad__q, 'b h s d -> b s h d') - return grad_q, grad_k, grad_v, grad_sinks - -- ++from .tilelang_kernel import sparse_mla_bwd, sparse_mla_fwd_interface + class AllGatherComm: """All gather communication with async operations""" - -@@ -131,212 +44,146 @@ class AllGatherComm: +@@ -131,212 +36,145 @@ class AllGatherComm: handle.wait() self.handles = [] @@ -146,9 +147,9 @@ index 89659a1d7..1def27c69 100644 '''Forward pass for the native attention function with context parallelism''' - # Assert einops exists - if not HAVE_EINOPS: - raise ImportError("einops is required by the attention CP but cannot be imported.") - +- if not HAVE_EINOPS: +- raise ImportError("einops is required by the attention CP but cannot be imported.") +- - # Initialize communication group and constants cp_size = 1 if pg is not None: @@ -164,8 +165,6 @@ index 89659a1d7..1def27c69 100644 - # Initialize KV buffers - kv_buffer = torch.empty( - (2, k.shape[0] * cp_size, k.shape[1], heads_k_stride, k.shape[3]), -+ s, b, heads, dim = q.shape -+ skv, _, kv_groups, _ = k.shape + + k_buffer = torch.empty( + (k.shape[0] * cp_size, k.shape[1], 1, k.shape[3]), @@ -232,9 +231,9 @@ index 89659a1d7..1def27c69 100644 + k_i = k_buffer + + s_, b_, h_, d_ = q_i.shape -+ q_i = einops.rearrange(q_i, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ q_i = q_i.transpose(0, 1).flatten().view(b_, s_, h_, d_) + s_, b_, h_, d_ = k_i.shape -+ k_i = einops.rearrange(k_i, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ k_i = k_i.transpose(0, 1).flatten().view(b_, s_, h_, d_) + zz_indices_i = zz_indices + b_, s_, g_, topk_ = zz_indices_i.shape + zz_indices_i = zz_indices_i.flatten().view(b_, s_, g_, topk_) @@ -245,7 +244,8 @@ index 89659a1d7..1def27c69 100644 + out_i, lse_i = sparse_mla_fwd_interface(q_i.contiguous(), k_i, zz_indices_i, zz_masks_i, dim_v, sm_scale = softmax_scale) + + # out: [B, seq_len_shard, h, dim] -> [seq_len, B, h, dim] -+ out_i = einops.rearrange(out_i, 'b s h d -> s b h d') ++ b_, s_, h_, d_ = out_i.shape ++ out_i = out_i.transpose(0, 1).flatten().view(s_, b_, h_, d_).contiguous() + + # outs: [[B, seq_len_shard, nheads // kv_group, dim], ...., [B, seq_len_shard, nheads // kv_group, dim]], repeat kv_group // heads_kv_stride times + # lses: [[B, seq_len_shard, heads_kv_stride], ...., [B, seq_len_shard, heads_kv_stride]], repeat kv_group // heads_kv_stride times @@ -353,13 +353,13 @@ index 89659a1d7..1def27c69 100644 + dk_list = [] + + s_, b_, h_, d_ = q.shape -+ q = einops.rearrange(q, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ q = q.transpose(0, 1).flatten().view(b_, s_, h_, d_) + s_, b_, h_, d_ = k_i.shape -+ k_i = einops.rearrange(k_i, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ k_i = k_i.transpose(0, 1).flatten().view(b_, s_, h_, d_) + s_, b_, h_, d_ = dout.shape -+ dout = einops.rearrange(dout, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ dout = dout.transpose(0, 1).flatten().view(b_, s_, h_, d_) + s_, b_, h_, d_ = out.shape -+ out = einops.rearrange(out, 's b h d -> b s h d').flatten().view(b_, s_, h_, d_) ++ out = out.transpose(0, 1).flatten().view(b_, s_, h_, d_) + b_, s_, h_ = lse.shape + lse = lse.flatten().view(b_, s_, h_) + zz_indices_i = zz_indices @@ -379,12 +379,16 @@ index 89659a1d7..1def27c69 100644 + + # TODO: needs casual = True, may not be compatible with zz + dq_i, _dk_i = sparse_mla_bwd(q_i, k_i, out_i, dout_i, zz_indices_i, zz_masks_i, lse_i, dim_v, sm_scale = softmax_scale) ++ ++ b_, s_, h_, d_ = dq_i.shape ++ dq_i = dq_i.transpose(0, 1).flatten().view(s_, b_, h_, d_).contiguous() ++ b_, s_, h_, d_ = _dk_i.shape ++ _dk_i = _dk_i.transpose(0, 1).flatten().view(s_, b_, h_, d_).contiguous() - # Rearrange gradients to (s, b, h, d) - dq_i = einops.rearrange(dq_i, 'b s h d -> s b h d') - _dk_i = einops.rearrange(_dk_i, 'b s h d -> s b h d') +- dq_i = einops.rearrange(dq_i, 'b s h d -> s b h d') +- _dk_i = einops.rearrange(_dk_i, 'b s h d -> s b h d') - _dv_i = einops.rearrange(_dv_i, 'b s h d -> s b h d') -+ if pg is None: dk_i = _dk_i - dv_i = _dv_i @@ -416,9 +420,11 @@ index 89659a1d7..1def27c69 100644 - dv = torch.cat(dv, dim=2) - return dq, dk, dv, None, None, None, None + dq = torch.cat(dq_list, dim=2) -+ dk = sum(dk_list) ++ dk_ = torch.cat(dk_list, dim=2) ++ dk = torch.sum(dk_, dim=2, keepdim=True) + + return dq, dk, None, None, None, None, None, None +\ No newline at end of file diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index fc994490b..b23d2e9a8 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py From d81f29ca286a6ebf7c03ad3c01f3ae8ca8113bf1 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Mon, 19 Jan 2026 23:32:19 -0800 Subject: [PATCH 27/30] update --- docker/deepseekv32/megatron.patch | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index c1f80e1e50d..886b73e90bc 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -1,5 +1,5 @@ diff --git a/megatron/core/transformer/dot_product_attention_context_parallel.py b/megatron/core/transformer/dot_product_attention_context_parallel.py -index 89659a1d7..2c1464fb6 100644 +index 89659a1d7..c69859a04 100644 --- a/megatron/core/transformer/dot_product_attention_context_parallel.py +++ b/megatron/core/transformer/dot_product_attention_context_parallel.py @@ -3,107 +3,12 @@ @@ -397,13 +397,14 @@ index 89659a1d7..2c1464fb6 100644 dk_i = torch.zeros( (k_i.shape[1] // cp_size, k_i.shape[0], k_i.shape[2], k_i.shape[3]), device=k_i.device, - dtype=k_i.dtype, - ) +- dtype=k_i.dtype, +- ) - dv_i = torch.zeros( - (v_i.shape[1] // cp_size, v_i.shape[0], v_i.shape[2], v_i.shape[3]), - device=v_i.device, - dtype=v_i.dtype, -- ) ++ dtype=torch.float32, + ) torch.distributed.reduce_scatter_tensor(dk_i, _dk_i, group=pg) - torch.distributed.reduce_scatter_tensor(dv_i, _dv_i, group=pg) @@ -421,7 +422,7 @@ index 89659a1d7..2c1464fb6 100644 - return dq, dk, dv, None, None, None, None + dq = torch.cat(dq_list, dim=2) + dk_ = torch.cat(dk_list, dim=2) -+ dk = torch.sum(dk_, dim=2, keepdim=True) ++ dk = torch.sum(dk_, dim=2, keepdim=True).to(torch.bfloat16) + + return dq, dk, None, None, None, None, None, None \ No newline at end of file @@ -1004,10 +1005,10 @@ index 000000000..d8f2425f0 +] diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py new file mode 100644 -index 000000000..b8ea416dd +index 000000000..83a259efa --- /dev/null +++ b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py -@@ -0,0 +1,274 @@ +@@ -0,0 +1,272 @@ +# ruff: noqa +import tilelang +from tilelang import language as T @@ -1273,22 +1274,20 @@ index 000000000..b8ea416dd + # Get kernels + preprocess_kernel = preprocess(B, S, H, D) + bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual) -+ postprocess_kernel = postprocess(B, S_kv, D, D_tail, kv_group) + + if delta is None: + delta = preprocess_kernel(o, do) + dkv = torch.zeros_like(kv, dtype=torch.float32) + dq = bwd_kernel(q, kv, do, indices, masks, lse, delta, dkv) -+ dkv = postprocess_kernel(dkv) + + return dq, dkv \ No newline at end of file diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py new file mode 100644 -index 000000000..d338a2fa6 +index 000000000..e247038de --- /dev/null +++ b/megatron/core/transformer/tilelang_kernel/sparse_mla_fwd.py -@@ -0,0 +1,190 @@ +@@ -0,0 +1,191 @@ +# ruff: noqa +import torch +import tilelang @@ -1409,6 +1408,7 @@ index 000000000..d338a2fa6 + for i_i in T.Pipelined(NI, num_stages=num_stages): + for bi_i in T.Parallel(BI): + mask[bi_i] = Masks[b_i, s_i, g_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i]] ++ + for bi_i, d_i in T.Parallel(BI, D): + KV_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] + for bi_i, d_i in T.Parallel(BI, D_tail): From e1e23050bfcbbd5ee584fa4b08f6ae2e9b34ca59 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Thu, 22 Jan 2026 13:58:34 -0800 Subject: [PATCH 28/30] update --- A.py | 1986 +++++++++++++++++++++++++++++ docker/deepseekv32/megatron.patch | 169 ++- 2 files changed, 2089 insertions(+), 66 deletions(-) create mode 100644 A.py diff --git a/A.py b/A.py new file mode 100644 index 00000000000..811f1c692d6 --- /dev/null +++ b/A.py @@ -0,0 +1,1986 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import warnings +from dataclasses import dataclass +from typing import Callable, List, Literal, Optional, Tuple, Union + +import torch +import torch.nn.functional as F + +from megatron.core.enums import Fp4Recipe, Fp8Recipe +from megatron.core.quantization.quant_config import RecipeConfig +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope +from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout + +from ..fusions.fused_bias_geglu import quick_gelu +from ..model_parallel_config import ModelParallelConfig +from ..utils import ( + get_te_version, + init_method_normal, + is_te_min_version, + is_torch_min_version, + scaled_init_method_normal, +) + +try: + from packaging.version import Version as PkgVersion + + HAVE_PACKAGING = True +except ImportError: + HAVE_PACKAGING = False + + +@dataclass +class TransformerConfig(ModelParallelConfig): + """Configuration object for megatron-core transformers. + + The initialization function has an argument for each parameter, + including those in ModelParallelConfig. + """ + + #################### + # model architecture + #################### + + num_layers: int = 0 + """Number of transformer layers in a transformer block.""" + + mtp_num_layers: Optional[int] = None + """Number of Multi-Token Prediction (MTP) Layers.""" + + mtp_loss_scaling_factor: Optional[float] = None + """Weighting factor of Multi-Token Prediction (MTP) loss.""" + + num_layers_in_first_pipeline_stage: Optional[int] = None + """Number of transformer layers on first pipeline stage. + None implies equal layer division across PP ranks.""" + + num_layers_in_last_pipeline_stage: Optional[int] = None + """Number of transformer layers on last pipeline stage. + None implies equal layer division across PP ranks.""" + + pipeline_model_parallel_layout: Optional[Union[str, list, PipelineParallelLayerLayout]] = None + """Custom definition of the pipeline parallel partitioning. + Support type: + - str: e.g., 'Et*3|(tt|)*29,m|L'. Stages are split by '|', replicated stages or layers + can be described with multiplication. Commas can be used cosmetically. + - list: e.g., [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]. + - PipelineParallelLayerLayout: a PipelineParallelLayerLayout object. + If given either a string or a list, it will be transferred into a PipelineParallelLayerLayout + in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers + in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., + vpp(i)pp(j), vpp(i)pp(j+1), ..., vpp(-1)pp(-2), vpp(-1)pp(-1). + In the inner lists of layers, 'embedding' or 'E' denotes the embedding layer, 'loss' or 'L' + denotes the loss function, and 'decoder' or 't' denotes the transformer decoder layer. + Examples: + [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]: + pp = 2, vpp = None + pp rank 0 holds: embedding, decoder + pp rank 1 holds: decoder*3, loss + 'E|(tt|)*2,(t|)*4,mL': + pp = 2, vpp = 4 + vpp rank 0 pp rank 0 holds: embedding + vpp rank 0 pp rank 1~2 holds: decoder*2 + vpp rank 0 pp rank 3 holds: decoder + vpp rank 1 pp rank 0~2 holds: decoder + vpp rank 1 pp rank 3 holds: mtp, loss""" + + account_for_embedding_in_pipeline_split: bool = False + """If set, the embedding layer will be treated as a standard transformer + layer in the context of partition and placement for pipeline parallelism.""" + + account_for_loss_in_pipeline_split: bool = False + """If set, the loss layer will be treated as a standard transformer + layer in the context of partition and placement for pipeline parallelism.""" + + hidden_size: int = 0 + """Transformer hidden size.""" + + num_attention_heads: int = 0 + """Number of transformer attention heads.""" + + attention_backend: AttnBackend = AttnBackend.auto + """Attention backend to run. By default we let transformer engine + decide the best backend to run (except in the case of local). + If attention backend is local we use the local pytorch implementation in mcore. + Users can specify exact backend by changing this config. """ + + softmax_scale: Optional[float] = None + """Softmax scale for attention scaling.""" + + softmax_type: Literal['vanilla', 'off-by-one', 'learnable'] = 'vanilla' + """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. + Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and + and learnable offset.""" + + num_query_groups: Optional[int] = None + """Number of query groups for group query attention. If None, normal attention is used.""" + + ffn_hidden_size: Optional[int] = None + """Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size + if not provided.""" + + kv_channels: Optional[int] = None + """Projection weights dimension in multi-head attention. This is set to hidden_size // + num_attention_heads if not provided.""" + + hidden_dropout: float = 0.1 + """Dropout probability for transformer hidden state.""" + + attention_dropout: float = 0.1 + """Post attention dropout probability.""" + + fp32_residual_connection: bool = False + """If true, move residual connections to fp32.""" + + # @jcasper should we keep this option? + apply_residual_connection_post_layernorm: bool = False + """If True, uses the original BERT residule connection ordering.""" + + layernorm_epsilon: float = 1e-5 + """Epsilon value for any LayerNorm operations.""" + + layernorm_zero_centered_gamma: bool = False + """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves + numerical stability.""" + + add_bias_linear: bool = True + """Include a bias term in all linear layers (QKV projections, after core attention, and two in + MLP layer).""" + + add_qkv_bias: bool = False + """Add a bias term only for QKV projections.""" + + gated_linear_unit: bool = False + """Use a gated linear unit for the first linear layer in the MLP.""" + + activation_func: Callable = F.gelu + """Activation function to use for the non-linearity in the MLP.""" + + activation_func_fp8_input_store: bool = False + """Store the input of MLP activation function in FP8 for backprop to save memory. + The stored input is casted back to the original precision before backprop compuatation.""" + + glu_linear_offset: float = 0.0 + """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only + used when gated_linear_unit is True""" + + activation_func_clamp_value: Optional[float] = None + """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func + is quick_gelu.""" + + num_moe_experts: Optional[int] = None + """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None + for no MoE.""" + + rotary_interleaved: bool = False + """True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of + first half and second half (LLaMa style). Default to False.""" + + window_size: Optional[Tuple[int, int]] = None + """If not None, then will use sliding window attention. The size of the window is specified by + the numbers inside the tuple; -1 is special value meaning "infinite window size".""" + + window_attn_skip_freq: Optional[Union[int, List[int]]] = None + """Frequency of full attention layers among sliding window attention layers. Accepts either: + - An integer N: Represents a (N-1):1 ratio, one full attention layer after (N-1) SWA layers. + - A list that defines a custom pattern, e.g.: [1,1,1,1,0,0,0,0], where 1 represents SWA. """ + + normalization: str = "LayerNorm" + """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`.""" + + qk_layernorm: bool = False + """Whether to apply `normalization` type of normalization to the query and key embeddings.""" + + qk_clip: bool = False + """Whether to clip the query and key weights. Needed for Muon MLA Model training.""" + + qk_clip_alpha: float = 0.5 + """The balancing alpha for qk-clip. Q = Q * (eta ** alpha)""" + + qk_clip_threshold: float = 100 + """The balancing threshold for qk-clip. eta = min(threshold / max_attention_logits, 1.0)""" + + log_max_attention_logit: bool = False + """Whether to log the max attention logit across whole model. Decoupled from qk_clip, + defualts to False. Setting qk_clip will automatically log the max logit""" + + attention_output_gate: bool = False + """Whether to apply output gate to the attention layers.""" + + test_mode: bool = False + """Whether to run real-time tests.""" + + calculate_per_token_loss: bool = False + """Whether cross entropy loss is calculated over the actual number of non-padded tokens in the + global batch, versus the default behavior of assuming all tokens are non-padded.""" + + multi_latent_attention: bool = False + """Whether to use multi-latent attention.""" + + no_rope_freq: Optional[Union[int, List[int]]] = None + """Controls which layers perform Rotary Position Embedding (RoPE). Accepts either: + An integer N: Creates a pattern where RoPE is skipped every N-1 layers. For example, + no_rope=4 means RoPE is applied for 3 layers, then skipped for 1 layer, repeating this pattern. + A list of integers: Defines a custom pattern where 1 means skip RoPE and 0 means apply RoPE. + For example, [0,1,1,0] means: apply RoPE, skip RoPE, skip RoPE, apply RoPE.""" + + moe_deepep_num_sms: int = 20 + """Number of SMs to use for DeepEP.""" + + moe_hybridep_num_sms: int = 16 + """Number of SMs to use for HybridEP. In pure NVL scenarios, + 16 SMs can generally achieve good bandwidth.""" + + #################### + # attention variant + #################### + experimental_attention_variant: Optional[str] = None + """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + + #################### + # attention variant: gated_delta_net + #################### + linear_attention_freq: Optional[Union[int, List[int]]] = None + """Frequency between LA (linear attention) layers + and SDPA (scaled dot-product attention) layers. + Accepts either: + - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer + - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]""" + + linear_conv_kernel_dim: Optional[int] = None + """Conv kernel dimension for the gated delta net.""" + + linear_key_head_dim: Optional[int] = None + """Query and key head dimension for the gated delta net.""" + + linear_value_head_dim: Optional[int] = None + """Value and gate head dimension for the gated delta net.""" + + linear_num_key_heads: Optional[int] = None + """Number of query and key heads for the gated delta net.""" + + linear_num_value_heads: Optional[int] = None + """Number of value and gate heads for the gated delta net.""" + + #################### + # attention variant: dsa + #################### + dsa_indexer_n_heads: Optional[int] = None + """Number of DSA indexer heads.""" + + dsa_indexer_head_dim: Optional[int] = None + """Dimension per DSA indexer head.""" + + dsa_indexer_topk: Optional[int] = None + """Number of top-k tokens to select in DSA indexer.""" + + dsa_indexer_loss_coeff: Optional[float] = None + """Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss.""" + + dsa_indexer_use_sparse_loss: Optional[bool] = None + """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the + top-k indices.""" + + #################### + # initialization + #################### + init_method: Optional[Callable] = None + """Method to initialize weights. Note that bias is always set to zero. Should be a function that + takes a single Tensor and initializes it. If None, will be set to + megatron.core.utils.init_method_normal(init_method_std) which is torch nn init normal with + mean=0.0 and std=init_method_std.""" + + output_layer_init_method: Optional[Callable] = None + """Method to initialize weights of the output layer of both attention and MLP blocks. If None, + will be set to megatron.core.utils.scaled_init_method_normal(init_method_std) which is torch nn + init normal with mean=0.0 and std=init_method_std / math.sqrt(2.0 * num_layers).""" + + init_method_std: float = 0.02 + """Standard deviation of the zero mean normal for the default initialization method, not used if + init_method and output_layer_init_method are provided.""" + + embedding_init_method: Optional[Callable] = None + """ + Method to initialize weights of the embedding layer. If None, will be set as described + in init_method above. + """ + + embedding_init_method_std: Optional[float] = None + """ + Standard deviation of the zero mean normal for the default initialization method for the + embedding layer. If None, will be set to init_method_std. + """ + + init_model_with_meta_device: bool = False + """ + If True, initializes the model with the meta device. This is helpful for + training of very large models. This feature is only works when megatron fsdp is turned on. + """ + + #################### + # mixed-precision + #################### + apply_query_key_layer_scaling: bool = False + """If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with + fp16.""" + + attention_softmax_in_fp32: bool = True + """If True, run attention masking and softmax in fp32. This should be True if + apply_query_key_layer_scaling is True.""" + + disable_bf16_reduced_precision_matmul: bool = False + """If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to + prevent matmul from using reduced precision accumulation when using BF16.""" + + #################### + # fusion + #################### + bias_activation_fusion: bool = False + """If True, fuses bias addition and the activation function when possible.""" + + masked_softmax_fusion: bool = False + """If True, uses softmax fusion.""" + + persist_layer_norm: bool = False + """If True, uses the persistent fused layer norm kernel. This kernel only supports a fixed set + of hidden sizes.""" + + memory_efficient_layer_norm: bool = False + """If True, and using local layers (not from TransformerEngine), tells Apex to use the memory + efficient fused LayerNorm kernel. Ignored if not using LayerNorm.""" + + bias_dropout_fusion: bool = False # TODO: this should be bias_dropout_add_fusion? + """If True, uses bias dropout fusion.""" + + apply_rope_fusion: bool = False + """If True, use fused RoPE kernel.""" + + use_fused_weighted_squared_relu: bool = False + """If True, uses fused weighted squared relu kernel when using MoE.""" + + fused_single_qkv_rope: bool = False + """If set, avoid splitting QKV before ROPE forward and avoid concatenating ROPE dgrads.""" + + #################### + # activation recomputation + #################### + recompute_granularity: Optional[str] = None + """Determines which type of activation recompute to use. Megatron-core supports 'selective' + activation checkpointing where the submodules set in --recompute-modules is checkpointed. + The default is "core_attn" which is the memory intensive part of attention. + These memory intensive activations are also less compute intensive which makes activation + checkpointing more efficient for LLMs (20B+). See Reducing Activation Recomputation in Large + Transformer Models (https://arxiv.org/abs/2205.05198) for more details. 'full' will checkpoint + the entire transformer layer. If None, no recompute is performed and all activations are saved. + If set, must be 'selective' or 'full'. 'selective' always uses all layers. + """ + + recompute_method: Optional[str] = None + """Determines which transformer layers will be recomputed. uniform will uniformly divide the + total number of transformer layers in a transformer block and recompute the input activation of + each divided chunk at the specified granularity. block will recompute the input activations for + only a set number of transformer layers per pipeline stage. The rest of the layers in the + pipeline stage will not have any activations recomputed. If None, and recompute is enabled, all + layers will do recomputation. If set, must be 'uniform' or 'block'.""" + + recompute_num_layers: Optional[int] = None + """When recompute_method is uniform, recompute_num_layers is the number of transformer layers in + each uniformly divided recompute unit. When recompute_method is block, recompute_num_layers is + the number of transformer layers to recompute within each pipeline stage. Must be None for + 'selective' activation checkpointing.""" + + distribute_saved_activations: Optional[bool] = None + """If True, distribute recomputed activations across the model parallel group.""" + + recompute_modules: Optional[List[str]] = None + """The submodules to recompute. + choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts". + default: ["core_attn"]. + "core_attn": recompute the core attention part of the transformer layer. + "moe_act": recompute the MoE MLP activation function. + "layernorm": recompute the input_layernorm and pre_mlp_layernorm. + "mla_up_proj": recompute the MLA up projection and RoPE applying parts. + "mlp": recompute the dense MLP submodule. + "moe": recompute the MoE layer. + "shared_experts": recompute the shared experts in the MoE layer. + "moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, + "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. + """ + + #################### + # fp8 related + #################### + fp8: Optional[str] = None + """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined + choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8 + activation and weight tensors and e5m2 for all FP8 output activation gradient tensors.""" + + fp8_recipe: Optional[str] = "delayed" + """If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined + choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed' + uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only, + 4) 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe.""" + + fp8_param: bool = False + """If set, keep the parameters in fp8 precision to save memory. This option must be used + together with fp8 mode (i.e., TransformerConfig.fp8 is not None). Note that not all parameters + will be converted to fp8; for example, biases will remain unchanged. The parameters affected are + primarily the weights of GEMMs. The specific parameters that will be converted to fp8 are + determined by TE.""" + + fp8_quantizer_factory: Optional[str] = None + """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. + Required when fp8_recipe is custom.""" + + fp8_margin: int = 0 + """Margin for the scaling factor computation.""" + + fp8_interval: int = 1 + """DEPRECATED from TransformerEngine v1.8.0. This flag is ignored. + Controls how often the scaling factor is recomputed. + """ + + fp8_amax_history_len: int = 1 + """The length of the amax history window used for scaling factor computation.""" + + fp8_amax_compute_algo: str = "most_recent" + """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2 + predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent` + always chooses the most recently seen value. + + """ + + fp8_wgrad: bool = True + """When set to False, override FP8 config options and do the wgrad computation + in higher precision.""" + + fp8_dot_product_attention: bool = False + """When set to True, use the FP8 implementation of Dot Product Attention.""" + + fp8_multi_head_attention: bool = False + """When set to True, use the FP8 implementation of Multi Head Attention.""" + + tp_only_amax_red: bool = False + """When set to True, reduce the FP8 AMAX only in the TP or TP-CP domain""" + + first_last_layers_bf16: bool = False + """If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8.""" + + num_layers_at_start_in_bf16: int = 1 + """Number of layers at the start of the model to keep in BF16 precision when + first_last_layers_bf16 is True.""" + + num_layers_at_end_in_bf16: int = 1 + """Number of layers at the end of the model to keep in BF16 precision when + first_last_layers_bf16 is True.""" + + use_kitchen: bool = False + """Use the kitchen extension for transformer quantization.""" + + #################### + # fp4 related + #################### + fp4: Optional[str] = None + """If set, enables the use of FP4 precision through Transformer Engine. Currently only + supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0).""" + + fp4_recipe: Optional[str] = "nvfp4" + """If set, enables the use of FP4 precision through Transformer Engine. Currently only + 'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture.""" + + fp4_param: bool = False + """If set, keep the parameters in fp4 precision to save memory. This option must be used + together with fp4 mode (i.e., TransformerConfig.fp4 is not None). Note that not all parameters + will be converted to fp4; for example, biases will remain unchanged.""" + + fp4_quantizer_factory: Optional[str] = None + """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. + Required when fp4_recipe is custom.""" + + #################### + # MoE related + #################### + moe_shared_expert_intermediate_size: Optional[int] = None + """Shared expert total ffn hidden size. + It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if + there are multiple shared experts. + None means no shared expert. + By default, the shared experts execute before the router. However, when + moe_shared_expert_overlap or overlap_moe_expert_parallel_comm is set, + the shared experts execute after the router, before the routed experts. + This makes the gradients from the router and the shared experts added in + different orders to the hidden_states, causing minor numerical differences + in the hidden_states gradient.""" + + moe_shared_expert_gate: bool = False + """Enable gate for shared expert.""" + + moe_shared_expert_overlap: bool = False + """Enable overlapping between shared expert computations and dispatcher communications. + Without this, the shared experts execute before the router.""" + + moe_layer_freq: Union[int, List[int]] = 1 + """Frequency between MoE layers and Dense layers. Accepts either: + - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers. + - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]""" + + moe_ffn_hidden_size: Optional[int] = None + """MoE Feed-Forward Network hidden size""" + + moe_router_load_balancing_type: Union[str, List[str]] = "aux_loss" + """The load balancing strategy for the router. + Options: + - "aux_loss": Load balancing loss used in GShard and SwitchTransformer, calculated at + micro-batch level. + - "seq_aux_loss": Load balancing loss used in DeepSeekV2 and DeepSeekV3, computes loss + for each individual sample. + - "global_aux_loss": Load balancing loss calculated at global batch level. + - "sinkhorn": Balancing algorithm used in S-BASE. + - "none": No load balancing. + A list of strings can be provided to combine multiple aux-loss load balancing types. + The default is "aux_loss". + """ + + moe_router_topk: int = 2 + """Number of experts to route to for each token.""" + + moe_router_topk_limited_devices: Optional[int] = None + """Number of EP ranks to consider for each token in group-limited routing, + DEPRECATED and replaced by moe_router_num_groups and moe_router_group_topk. + """ + + moe_router_padding_for_quantization: Optional[bool] = False + """Whether to pad the routing_map to make sure the number of tokens each expert receives + is a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the explicit + padding in the GroupedMLP layer.""" + + moe_router_padding_for_fp8: Optional[bool] = False + """[Compatibility alias for moe_router_padding_for_quantization] + Enabling this will also enable moe_router_padding_for_quantization.""" + + moe_router_num_groups: Optional[int] = None + """Number of groups to divide experts into for group-limited routing. + When using group-limited routing: + 1. Experts are divided into 'moe_router_num_groups' equal-sized groups + 2. For each token, 'moe_router_group_topk' groups are selected based on sum of + top-('moe_router_topk'/'moe_router_group_topk') routing scores within each group + 3. From these selected groups, 'moe_router_topk' individual experts are chosen + Two common use cases: + - Device-limited routing: Set 'moe_router_num_groups' equal to expert parallel size (EP) + to limit each token to experts on a subset of devices + (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) + - Node-limited routing: Set 'moe_router_num_groups' equal to number of nodes in EP group + to limit each token to experts on a subset of nodes + (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437) + """ + + moe_router_group_topk: Optional[int] = None + """Number of selected groups for group-limited routing.""" + + moe_router_pre_softmax: bool = False + """Enable pre-softmax(pre-sigmoid) routing for MoE, which means softmax is before the + top-k selection. + By default, softmax is done after top-k.""" + + moe_router_topk_scaling_factor: Optional[float] = None + """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax + enabled. Defaults to None, which means no scaling.""" + + moe_router_score_function: str = "softmax" + """Score function for MoE routing. Can be "softmax" or "sigmoid".""" + + moe_router_dtype: Optional[str] = None + """Data type for routing and expert output weighted averaging. Using fp32 or fp64 can + improve stability especially when the number of experts is large (e.g. finegrained-moe). + None means no changes for dtype.""" + + moe_router_enable_expert_bias: bool = False + """TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. + The routing decision is based on the sum of the routing scores and the expert bias. + See https://arxiv.org/abs/2408.15664 for details.""" + + moe_router_bias_update_rate: float = 1e-3 + """The expert bias is updated based on the number of assigned tokens to each expert + in a global batch, where the bias is increased for the experts with less assigned tokens + and decreased for the experts with more assigned tokens. + The default value 1e-3 is same as that used in DeepSeekV3.""" + + moe_router_force_load_balancing: bool = False + """[Experimental] Force load balancing with random logits for MoE router, supports naive topk + and group-limited topk. This is an experimental feature and only for benchmark.""" + + moe_grouped_gemm: bool = False + """When there are multiple experts per rank, compress multiple local (potentially small) gemms + in a single kernel launch to improve the utilization and performance by leveraging the Grouped + GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). + """ + + moe_use_legacy_grouped_gemm: bool = False + """Use legacy GroupedMLP rather than TEGroupedMLP. + Note: The legacy one will be deprecated soon.""" + + moe_aux_loss_coeff: Union[float, List[float]] = 0.0 + """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. + If a list of load balancing types is provided for `moe_router_load_balancing_type`, + a corresponding list of coefficients should be provided here.""" + + moe_z_loss_coeff: Optional[float] = None # 1e-3 would be a good start value for z-loss + """Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended.""" + + moe_input_jitter_eps: Optional[float] = None + """Add noise to the input tensor by applying jitter with a specified epsilon value.""" + + moe_token_dropping: bool = False + """This feature involves selectively dropping and padding tokens for each expert to achieve a + specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is + currently unsupported so should remain False.""" + + moe_token_dispatcher_type: str = "allgather" + """The type of token dispatcher to use. The default is 'allgather'. + Options are 'allgather','alltoall' and 'flex'.""" + + moe_enable_deepep: bool = False + """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models.""" + + moe_flex_dispatcher_backend: str = "deepep" + """[Experimental] The backend to use for flex token dispatcher. The default is "deepep". + Options are "deepep" and "hybridep". Currently only "hybridep" backend supports + the MNNVL case.""" + + moe_per_layer_logging: bool = False + """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.""" + + moe_expert_capacity_factor: Optional[float] = None + """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token + will be dropped. The default is None.""" + + moe_pad_expert_input_to_capacity: bool = False + """moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match + the expert capacity length, effective only after the moe_expert_capacity_factor is set. The + default setting is False.""" + + moe_token_drop_policy: str = "probs" + """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with + the lowest probabilities will be dropped. If "position", tokens at the end of each batch will + be dropped. + """ + + moe_layer_recompute: bool = False + """Memory optimization: checkpointing moe_layer to save actiavtion memory.""" + + moe_permute_fusion: bool = False + """Fuse token rearrangement ops during token dispatching.""" + + moe_router_fusion: bool = False + """Fuse ops in routing and aux loss calculation.""" + + moe_apply_probs_on_input: bool = False + """Apply probs on input of experts instead of applying after activation and glu.""" + + ################## + # Context Parallel + ################## + cp_comm_type: Optional[Union[str, List[str]]] = None + """Inter-gpu communication type for context parallelism. + str: all layers share same communication type. + List[str]: each layer has its separate communication type. + cp_comm_type of each layer can be "p2p" or "all_gather" or "a2a" or "a2a+p2p". + "p2p": Exchange KV chunks with P2P communications in ring topology. P2P is async and can be + overlapped with attention compute. + "all_gather": All-gather to get full sequence of KV before attention. The all-gather is not + async, and cannot be overlapped. + "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP group, and gather to get + full sequence of QKV. + "a2a+p2p": A hierarchical implementation of context parallelism to attention. + It uses A2A communications in low-level CP groups (e.g., via NVLink), + and P2P communications in high-level CP groups (e.g., via IBLink). + """ + + ################## + # Cuda Graphs + ################## + enable_cuda_graph: bool = False + """DEPRECATED and replaced by cuda_graph_impl. + When set to true, either partial CUDA graph (1/many CUDA graph per layer) or full iteration + CUDA graph (1 CUDA graph for whole iteration excluding optimizer) is enabled. --cuda-graph-scope + determines the scope of graph capture.""" + + cuda_graph_use_single_mempool: bool = False + """When set to true, cudagraphs will be captured inside a single mempool, in which all + cudagraphs may only be used once per step. If false, cudagraphs may be reused across + microbatches. Enabling may reduce cudagraph memory overheads due to memory fragmentation, + however may greatly increase the number of cudagraphs created when the number of microbatches + is high.""" + + cuda_graph_retain_backward_graph: bool = False + """When set to true, cudagraph backward passes will be graph captured with 'retain_grad=True' + This may enable cudagraphs for certain modules that are not completely cudagraph safe. For + more details, see: https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html.""" + + cuda_graph_warmup_steps: int = 3 + """Number of warmup steps for CUDA graphs""" + + external_cuda_graph: bool = False + """DEPRECATED and replaced by cuda_graph_impl. + When set to true, TransformerLayer layers are swapped with user provided CUDA graphs.""" + + cuda_graph_impl: str = "none" + """Determines the CUDA graph capture implementation. + "none": no CUDA graph. + "local": capture the CUDA graph using MCore local implementation. Either partial CUDA graph + (1/many CUDA graph per layer) or full iteration CUDA graph (1 CUDA graph for whole iteration + excluding optimizer) is enabled. + "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" + + cuda_graph_scope: Optional[List[CudaGraphScope]] = None + """Determines the CUDA graphs capturing scope. + When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", + "moe_router", "moe_preprocess", "mamba". None means the full layer. + When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope + to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" + + #################### + # miscellaneous + #################### + clone_scatter_output_in_embedding: bool = True + """When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer + to facilitate garbage collection of input.""" + + disable_parameter_transpose_cache: bool = False + """When set to true, the parameter transposes are not cached for subsequent iterations.""" + + config_logger_dir: str = "" + """When non-empty, dumps entry-point configs to config_logger_dir""" + + flash_decode: bool = False + """ Use the optimized flash decoding kernel during inference. """ + + use_te_activation_func: bool = False + """Whether to use ffn activation functions implemented by TransformerEngine""" + + use_te_rng_tracker: bool = False + """ Whether to use the TE or MCore version of the RNG tracker. """ + + inference_rng_tracker: bool = False + """ Whether we should instantiate a separate RNG tracker for inference. """ + + inference_sampling_seed: int = 42 + """ Random seed to use for sampling during inference. """ + + symmetric_ar_type: Optional[str] = None + """Type of symmetric all reduce to use""" + + mrope_section: Optional[List[int]] = None + """ Multimodal rope section is for channel dimension of temporal, height and width + in rope calculation. """ + + is_hybrid_model: bool = False + """ Indicates whether this is a hybrid model. """ + + mamba_state_dim: int = 128 + """The dimensionality of the state representation in Mamba layers.""" + + mamba_head_dim: int = 64 + """The dimensionality of the heads in the Mamba layers.""" + + mamba_num_groups: int = 8 + """The number of groups used in Mamba layers.""" + + mamba_num_heads: Optional[int] = None + """The number of heads used in Mamba layers. + If None, the number of heads will be hidden_size * expand // mamba_head_dim.""" + + use_mamba_mem_eff_path: bool = True + """If True, use the memory efficient path for Mamba layers.""" + + mlp_chunks_for_prefill: int = 1 + """The number of chunks along the sequence dimension to use for MLP computation + during prefill.""" + + heterogeneous_block_specs: bool = False + """Whether to use heterogeneous block specs (nemotron-nas architecture).""" + + hetereogenous_dist_checkpoint: bool = False + """Whether to use heterogenous layers in distributed checkpoint.""" + + #################### + # Quantization + #################### + quant_recipe: Optional[RecipeConfig] = None + """Configuration of any quantization to be applied to the model""" + + transformer_impl: str = "transformer_engine" + """Transformer implementation to use. + Options are 'transformer_engine' for Transformer Engine and 'local' for MCore.""" + + fallback_to_eager_attn: bool = False + """Whether to fallback to eager attention in TE implementation. + Suggested for when desired features are not available in TE implementation.""" + + ##################################### + # Fine-grained Activation Offloading + ##################################### + fine_grained_activation_offloading: bool = False + """If True, offload the input of the specified modules to the CPU. + Fine-grained activation offloading is a module-level offloading method + instead of a layer-level offloading method like cpu_offloading.""" + + offload_modules: Optional[list[str]] = None + """The submodules to offload its input. + choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", + "mlp_norm", "expert_fc1", "moe_act". + "attn_norm": offload the input of the normalization in the attention part. + "qkv_linear": offload the input of the qkv linear part. + "core_attn": offload the input of the core attention part. + "attn_proj": offload the input of the attn linear projection part. + "mlp_norm": offload the input of the normalization in the mlp part. + "expert_fc1": offload the input of the expert fc1 part. + "moe_act": offload the input of the moe act part. + """ + min_offloaded_tensor_size: int = 1024 * 1024 + """The minimum size of the tensor to be offloaded.""" + + def __post_init__(self): + """Python dataclass method that is used to modify attributes after initialization. + See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more + details. + """ + super().__post_init__() + if self.fp16 and self.bf16: + raise ValueError( + f"Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True." + ) + + # Apply BF16 matmul precision setting if needed + if self.bf16 and self.disable_bf16_reduced_precision_matmul: + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + + if self.num_attention_heads % self.tensor_model_parallel_size != 0: + raise ValueError( + f"num_attention_heads ({self.num_attention_heads}) must be a multiple of " + f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." + ) + + if self.ffn_hidden_size is None: + self.ffn_hidden_size = 4 * self.hidden_size + + if self.kv_channels is None: + self.kv_channels = self.hidden_size // self.num_attention_heads + + if self.num_query_groups is None: + self.num_query_groups = self.num_attention_heads + + if self.num_query_groups % self.tensor_model_parallel_size != 0: + raise ValueError( + f"num_query_groups ({self.num_query_groups}) must be a multiple of " + f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." + ) + + if self.experimental_attention_variant in ["gated_delta_net"]: + assert ( + self.linear_attention_freq is not None + ), f"linear_attention_freq must be set for linear attention." + + if self.experimental_attention_variant == "gated_delta_net": + # Check required parameters + assert ( + self.linear_conv_kernel_dim is not None + ), "linear_conv_kernel_dim must be set for gated delta net." + assert ( + self.linear_key_head_dim is not None + ), "linear_key_head_dim must be set for gated delta net." + assert ( + self.linear_value_head_dim is not None + ), "linear_value_head_dim must be set for gated delta net." + assert ( + self.linear_num_key_heads is not None + ), "linear_num_key_heads must be set for gated delta net." + assert ( + self.linear_num_value_heads is not None + ), "linear_num_value_heads must be set for gated delta net." + assert self.linear_num_value_heads % self.linear_num_key_heads == 0, ( + f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of " + f"linear_num_key_heads ({self.linear_num_key_heads})." + ) + + # Check tensor parallelism compatibility + assert ( + self.linear_num_key_heads % self.tensor_model_parallel_size == 0 + ), "linear_num_key_heads must be a multiple of tensor_model_parallel_size." + assert ( + self.linear_num_value_heads % self.tensor_model_parallel_size == 0 + ), "linear_num_value_heads must be a multiple of tensor_model_parallel_size." + + # Do not support yet, but coming soon. + assert self.context_parallel_size == 1, ( + f"Gated delta net does not support context parallel for now," + f" but got {self.context_parallel_size=}." + ) + elif self.experimental_attention_variant == "dsa": + # assert ( + # self.context_parallel_size == 1 + # ), "Currently context parallelism is not supported by DSAttention!" + assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" + + if self.fp8: + # cannot support first last layer bf16 with delayed scaling + if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed: + raise ValueError("Delayed scaling does not support first / last layer in BF16.") + + # max bf16 layers per pipeline stage + max_bf16_layers_per_pipeline_stage = ( + self.num_layers // self.pipeline_model_parallel_size + ) + + # check start/end bf16 layer counts are valid + if self.first_last_layers_bf16: + if ( + self.num_layers_at_start_in_bf16 < 0 + or self.num_layers_at_start_in_bf16 > max_bf16_layers_per_pipeline_stage + ): + raise ValueError( + f"num_layers_at_start_in_bf16 ({self.num_layers_at_start_in_bf16}) must be " + f"between 0 and number of layers per pipeline stage " + f"({max_bf16_layers_per_pipeline_stage})." + ) + if ( + self.num_layers_at_end_in_bf16 < 0 + or self.num_layers_at_end_in_bf16 > max_bf16_layers_per_pipeline_stage + ): + raise ValueError( + f"num_layers_at_end_in_bf16 ({self.num_layers_at_end_in_bf16}) must be " + f"between 0 and number of layers per pipeline stage " + f"({max_bf16_layers_per_pipeline_stage})." + ) + + if self.fp8_recipe == Fp8Recipe.custom: + if not self.fp8_quantizer_factory: + raise ValueError( + "fp8_quantizer_factory must be provided when fp8_recipe is 'custom'. " + "Specify a Python import path (e.g., package.module.quantizer_factory) " + "via --fp8-quantizer-factory." + ) + + if self.fp8_param and not self.fp8: + raise ValueError("fp8_param must be used together with fp8 mode.") + + # FP4 validation + if self.fp4_param and not self.fp4: + raise ValueError("fp4_param must be used together with fp4 mode.") + + if self.fp4 and self.fp8: + raise ValueError("fp4 and fp8 cannot be used simultaneously. Please choose one.") + + if self.fp4 and self.fp4_recipe == Fp4Recipe.custom: + if not self.fp4_quantizer_factory: + raise ValueError( + "fp4_quantizer_factory must be provided when fp4_recipe is 'custom'. " + "Specify a Python import path (e.g., package.module.quantizer_factory) " + "via --fp4-quantizer-factory." + ) + + if self.apply_query_key_layer_scaling: + self.attention_softmax_in_fp32 = True + + if self.expert_model_parallel_size > 1 and self.num_moe_experts is None: + raise ValueError("num_moe_experts must be non None to use expert-parallel.") + + if self.num_moe_experts is not None and self.num_moe_experts <= 0: + raise ValueError("num_moe_experts must be non-negative.") + + if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None: + self.moe_ffn_hidden_size = self.ffn_hidden_size + warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.") + + if self.num_moe_experts is None: + assert ( + self.moe_ffn_hidden_size is None + ), "moe_ffn_hidden_size must be None when num_experts is not set." + + if self.moe_enable_deepep: + if self.moe_token_dispatcher_type != "flex": + raise ValueError("DeepEP backend is only supported with flex token dispatcher.") + if self.moe_flex_dispatcher_backend == "hybridep": + raise ValueError("Only one backend is supported for flex token dispatcher.") + self.moe_flex_dispatcher_backend = "deepep" + warnings.warn( + "moe_enable_deepep is deprecated." + "Please use --moe-flex-dispatcher-backend=deepep instead." + ) + + if self.moe_token_dispatcher_type == "flex": + if self.moe_pad_expert_input_to_capacity and ( + self.moe_enable_deepep or self.moe_flex_dispatcher_backend == "deepep" + ): + raise ValueError( + "Flex token dispatcher with deepep backend does not support " + "moe_pad_expert_input_to_capacity" + ) + + if self.moe_shared_expert_intermediate_size is not None: + if self.moe_shared_expert_intermediate_size <= 0: + raise ValueError( + f"moe_shared_expert_intermediate_size must be " + f"num_shared_experts * ffn_size_of_each_shared_expert, " + f"but got {self.moe_shared_expert_intermediate_size}" + ) + if self.moe_shared_expert_overlap and self.moe_token_dispatcher_type not in [ + "alltoall" + ]: + raise ValueError( + f"moe_shared_expert_overlap only works with alltoall token dispatcher." + ) + + if isinstance(self.moe_router_load_balancing_type, list): + assert isinstance(self.moe_aux_loss_coeff, list) and len( + self.moe_aux_loss_coeff + ) == len(self.moe_router_load_balancing_type), ( + "moe_aux_loss_coeff must be a list of the same length as " + "moe_router_load_balancing_type" + ) + + if self.moe_expert_capacity_factor is not None: + if self.moe_expert_capacity_factor < 0: + self.moe_expert_capacity_factor = None + if isinstance(self.moe_router_load_balancing_type, list): + for load_balancing_type in self.moe_router_load_balancing_type: + if load_balancing_type not in [ + "aux_loss", + "seq_aux_loss", + "global_aux_loss", + "none", + ]: + raise ValueError( + "moe_expert_capacity_factor only works with aux_loss, " + "seq_aux_loss, global_aux_loss or none load balancing" + ) + elif self.moe_router_load_balancing_type not in [ + "aux_loss", + "seq_aux_loss", + "global_aux_loss", + "none", + ]: + raise ValueError( + "moe_expert_capacity_factor only works with aux_loss, " + "seq_aux_loss, global_aux_loss or none load balancing" + ) + + if self.moe_pad_expert_input_to_capacity: + if self.moe_expert_capacity_factor is None: + raise ValueError( + "moe_expert_capacity_factor must be set to use moe_pad_expert_input_to_capacity" + ) + + if self.cpu_offloading and ( + self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers + ): + raise ValueError( + f"CPU offloading can be done only for layers less than {self.num_layers}" + ) + + if self.cpu_offloading and self.pipeline_model_parallel_size > 1: + raise ValueError( + "Currently there is no support for Pipeline parallelism with CPU offloading" + ) + + if self.cpu_offloading and self.recompute_granularity is not None: + raise ValueError( + "CPU offloading does not work when activation recomputation is enabled" + ) + + if self.recompute_granularity is not None: + if self.recompute_granularity not in ["full", "selective"]: + raise ValueError( + f'When using recompute_granuarlity: {self.recompute_granularity} must be "full"' + 'or "selective".' + ) + + if self.recompute_method is not None: + if self.recompute_method not in ["block", "uniform"]: + raise ValueError( + f'recompute_method: {self.recompute_method} must be "block" or "uniform".' + ) + elif self.recompute_granularity != "selective": + raise ValueError( + f"Using recompute_granularity: {self.recompute_granularity} so " + 'recompute_method must be "block" or "uniform"' + ) + + if self.recompute_granularity != "selective" and self.recompute_num_layers is None: + raise ValueError( + f"When using recompute_granularity: {self.recompute_granularity} " + "recompute_num_layers must be between " + "1 and num_layers_per_pipeline_rank: " + f"{self.num_layers // self.pipeline_model_parallel_size}" + ) + elif ( + self.recompute_granularity == "selective" and self.recompute_num_layers is not None + ): + raise ValueError( + f"When using recompute_granularity: {self.recompute_granularity} " + "recompute_num_layers must be None." + ) + + if self.distribute_saved_activations and self.sequence_parallel: + raise ValueError( + f"distribute_saved_activations: {self.distribute_saved_activations} must be " + f"false when sequence parallel is enabled: {self.sequence_parallel}" + ) + + if self.recompute_modules is None: + self.recompute_modules = ["core_attn"] + + if self.recompute_granularity == "selective": + if len(self.recompute_modules) > 0: + allowed_modules = { + "core_attn", + "moe_act", + "layernorm", + "mla_up_proj", + "mlp", + "moe", + "shared_experts", + } + invalid_modules = set(self.recompute_modules) - allowed_modules + assert not invalid_modules, ( + f"Invalid choices for recompute_modules: {invalid_modules}. " + f"Allowed modules are: {allowed_modules}" + ) + + if "moe_act" in self.recompute_modules and not self.moe_grouped_gemm: + raise ValueError( + "moe_act in recompute_modules is only supported with moe_grouped_gemm." + ) + + if "mla_up_proj" in self.recompute_modules and not self.multi_latent_attention: + raise ValueError( + "mla_up_proj in recompute_modules is only supported with " + "multi_latent_attention." + ) + + if "core_attn" in self.recompute_modules: + warnings.warn( + "If you are using transformer_engine as the transformer implementation, " + "the core_attn is from transformer_engine and may be the fused version. " + "For fused attention, you have no need to set 'core_attn' to recompute. " + "Please check that the core_attn recompute is really needed." + ) + + if "shared_experts" in self.recompute_modules: + if ( + self.moe_shared_expert_intermediate_size is not None + and self.moe_shared_expert_overlap + ): + raise ValueError( + "shared_experts recompute cannot work with --moe-shared-expert-overlap." + ) + + if self.fp8: + if "moe_act" in self.recompute_modules or "layernorm" in self.recompute_modules: + if self.fp8_recipe == 'delayed': + raise ValueError( + "Delayed scaling does not support moe_act and layernorm recompute " + "for fp8." + ) + if not is_te_min_version("2.6.0dev0"): + raise ValueError( + "moe_act and layernorm recompute for fp8 needs " + "transformer-engine>=2.6.0dev0, " + f"but your version is {get_te_version()}." + ) + + if self.moe_layer_recompute: + warnings.warn( + "--moe-layer-recompute is deprecated. " + "Use --recompute-granularity selective --recompute-modules moe_layer instead." + ) + if self.recompute_granularity == "full": + raise ValueError( + "Do not set --moe-layer-recompute with full recompute granularity. " + ) + self.recompute_granularity = "selective" + if "moe" not in self.recompute_modules: + self.recompute_modules.append("moe") + + if self.fine_grained_activation_offloading: + assert ( + not self.cpu_offloading + ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading." + assert self.offload_modules is not None and len(self.offload_modules) > 0 + allowed_modules = { + "core_attn", + "attn_proj", + "expert_fc1", + "moe_act", + "attn_norm", + "mlp_norm", + "qkv_linear", + } + invalid_modules = set(self.offload_modules) - allowed_modules + assert not invalid_modules, ( + f'Invalid choices for offload_modules: {invalid_modules}. ' + f'Allowed modules are: {allowed_modules}' + ) + if "attn_proj" in self.offload_modules and "core_attn" not in self.offload_modules: + raise ValueError( + "attn_proj cannot be set to offload_modules alone without core_attn " + "because the input of attn_proj is the output of core_attn, " + "which is needed in core_attn.backward()." + ) + + if ( + self.num_layers_in_first_pipeline_stage is not None + or self.num_layers_in_last_pipeline_stage is not None + ) and ( + self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split + ): + raise ValueError( + "num_layers_in_first_pipeline_stage and num_layers_in_last_pipeline_stage cannot be" + "set at the same time with account_for_embedding_in_pipeline_split" + "and account_for_loss_in_pipeline_split" + ) + + # PP layout + if self.pipeline_model_parallel_layout is not None: + # If pipeline layout is set, we will check the conflicts + # with other pipeline layout arguments. + any_conflict = ( + self.num_layers_in_first_pipeline_stage is not None + or self.num_layers_in_last_pipeline_stage is not None + or self.account_for_embedding_in_pipeline_split + or self.account_for_loss_in_pipeline_split + ) + if any_conflict: + raise ValueError( + "pipeline_model_parallel_layout cannot be set" + " with other pipeline layout arguments." + f" {self.num_layers_in_first_pipeline_stage=}," + f" {self.num_layers_in_last_pipeline_stage=}," + f" {self.account_for_embedding_in_pipeline_split=}," + f" {self.account_for_loss_in_pipeline_split=}." + ) + + # Transfer pipeline_model_parallel_layout from str or list to + # PipelineParallelLayerLayout + if isinstance(self.pipeline_model_parallel_layout, str): + self.pipeline_model_parallel_layout = PipelineParallelLayerLayout.from_str( + layout=self.pipeline_model_parallel_layout, + pipeline_model_parallel_size=self.pipeline_model_parallel_size, + ) + elif isinstance(self.pipeline_model_parallel_layout, list): + # Since list is not hashable, the initialization will not be cached. + self.pipeline_model_parallel_layout = PipelineParallelLayerLayout( + layout=self.pipeline_model_parallel_layout, + pipeline_model_parallel_size=self.pipeline_model_parallel_size, + ) + + # Check whether the input VPP size conflicts with the PP layout + detected_vpp_size = ( + self.pipeline_model_parallel_layout.virtual_pipeline_model_parallel_size + ) + if self.virtual_pipeline_model_parallel_size is not None: + assert self.virtual_pipeline_model_parallel_size == detected_vpp_size, ( + f"virtual_pipeline_model_parallel_size conflicts with" + f" pipeline_model_parallel_layout," + f" ({self.virtual_pipeline_model_parallel_size=}, " + f" {detected_vpp_size=})" + ) + elif detected_vpp_size > 1: + self.virtual_pipeline_model_parallel_size = detected_vpp_size + + # Check whether the layout is valid. + self.mtp_standalone = self.pipeline_model_parallel_layout.validate_layer_layout( + num_layers=self.num_layers, mtp_num_layers=self.mtp_num_layers + ) + + # Uneven PP + elif ( + self.num_layers_in_first_pipeline_stage is not None + or self.num_layers_in_last_pipeline_stage is not None + ): + pipeline_parallel_size = self.pipeline_model_parallel_size + num_layers = self.num_layers + + if self.num_layers_in_first_pipeline_stage is not None: + if self.num_layers_in_first_pipeline_stage <= 0: + raise ValueError("num_layers_in_first_pipeline_stage must be larger than 0") + + if self.virtual_pipeline_model_parallel_size is not None: + if ( + self.num_layers_in_first_pipeline_stage + % self.virtual_pipeline_model_parallel_size + != 0 + ): + raise ValueError( + f"number of layers at first stage: " + f"{self.num_layers_in_first_pipeline_stage}" + f"must be divisible by virtual pipeline" + f"parallel degree {self.virtual_pipeline_model_parallel_size}" + ) + num_layers -= self.num_layers_in_first_pipeline_stage + pipeline_parallel_size -= 1 + + if self.num_layers_in_last_pipeline_stage is not None: + if self.num_layers_in_last_pipeline_stage <= 0: + raise ValueError("num_layers_in_last_pipeline_stage must be larger than 0") + + if self.virtual_pipeline_model_parallel_size is not None: + if ( + self.num_layers_in_last_pipeline_stage + % self.virtual_pipeline_model_parallel_size + != 0 + ): + raise ValueError( + f"number of layers at last stage: " + f"{self.num_layers_in_last_pipeline_stage}" + f"must be divisible by virtual pipeline" + f"parallel degree {self.virtual_pipeline_model_parallel_size}" + ) + num_layers -= self.num_layers_in_last_pipeline_stage + pipeline_parallel_size -= 1 + + # Here pipeline_parallel_size is the number of middle PP stages. If there are middle + # PP stages, check number of layers at middle stage is divisible by middle PP size. + if pipeline_parallel_size and not num_layers % pipeline_parallel_size == 0: + raise ValueError( + f"number of layers at middle stage: {num_layers} must be divisible by" + f"the middle pipeline model parallel size {pipeline_parallel_size}" + ) + + # If there are middle PP stages, check number of layers + # on each middle PP rank is divisible by VPP size. + if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None: + num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size + if ( + not num_layers_per_middle_pipeline_rank + % self.virtual_pipeline_model_parallel_size + == 0 + ): + raise ValueError( + f"number of layers on each middle pipeline rank:" + f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual" + f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}" + ) + + elif ( + self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split + ): + if self.virtual_pipeline_model_parallel_size is None: + num_layers = self.num_layers + + if self.account_for_embedding_in_pipeline_split: + num_layers += 1 + + if self.account_for_loss_in_pipeline_split: + num_layers += 1 + + if not num_layers % self.pipeline_model_parallel_size == 0: + raise ValueError( + f"number of middle layers: {num_layers} must be divisible by " + f"middle pipeline_model_parallel_size {self.pipeline_model_parallel_size}" + ) + else: + num_layers = self.num_layers + if self.account_for_embedding_in_pipeline_split: + num_layers += 1 + + if self.account_for_loss_in_pipeline_split: + num_layers += 1 + + if not num_layers % self.pipeline_model_parallel_size == 0: + raise ValueError( + f"num_layers: {num_layers} after enable" + f"account_for_embedding_in_pipeline_split or " + f"account_for_loss_in_pipeline_split must be divisible" + f"by pipeline_model_parallel_size " + f"{self.pipeline_model_parallel_size}" + ) + + num_layers_per_pipeline_rank = num_layers // self.pipeline_model_parallel_size + if ( + not num_layers_per_pipeline_rank % self.virtual_pipeline_model_parallel_size + == 0 + ): + raise ValueError( + f"number of layers on each pipeline rank: {num_layers_per_pipeline_rank}" + f"(after enable account_for_embedding_in_pipeline_split or " + f"account_for_loss_in_pipeline_split) must be divisible by" + f"virtual_pipeline_model_parallel_size" + f"{self.virtual_pipeline_model_parallel_size}" + ) + + if self.apply_query_key_layer_scaling: + self.attention_softmax_in_fp32 = True + + if self.bias_activation_fusion: + if self.activation_func not in [F.gelu, F.silu, quick_gelu]: + raise ValueError( + "When bias_activation_fusion is True, activation function should be either " + "gelu, swiglu, or quick_geglu" + ) + if ( + self.activation_func == F.gelu + and not self.gated_linear_unit + and not self.add_bias_linear + ): + raise ValueError( + "When bias_activation_fusion is True, gated_linear_unit is False " + "and activation function is gelu, add_bias_linear must also be True." + ) + if self.activation_func == quick_gelu and not self.gated_linear_unit: + raise ValueError( + "When bias_activation_fusion is True and activation function is quick_gelu, " + "gated_linear_unit must be True." + ) + if self.glu_linear_offset != 0.0 and self.activation_func != quick_gelu: + raise ValueError( + "When bias_activation_fusion is True and glu_linear_offset is non-zero, " + "activation function must be quick_gelu." + ) + + if self.use_te_activation_func: + raise ValueError( + "bias_activation_fusion and use_te_activation_func cannot be both true. " + "If you use bias in MLP FC1, we recommend setting bias_activation_fusion " + "to True and use_te_activation_func to False." + ) + + if self.use_te_activation_func: + if self.activation_func not in (F.gelu, F.silu, F.relu): + raise ValueError( + "TransformerEngine only support gelu, geglu, silu, swiglu, relu, reglu. " + "If you don't want to use TransformerEngine activation function, set " + "use_te_activation_func to False" + ) + + if self.activation_func_fp8_input_store: + if self.activation_func != F.silu or not self.gated_linear_unit: + raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.") + + if self.apply_rope_fusion: + if self.multi_latent_attention: + warnings.warn( + "apply_rope_fusion for multi-latent attention only supports training. " + "It is experimental and may change in future versions." + ) + else: + if self.rotary_interleaved: + if not is_te_min_version("2.3.0"): + raise ValueError( + "rotary_interleaved does not work with apply_rope_fusion for " + "TE < 2.3.0. Please install TE >= 2.3.0" + ) + + from megatron.core.models.common.embeddings.rope_utils import ( + fused_apply_rotary_pos_emb, + fused_apply_rotary_pos_emb_thd, + ) + + if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None: + raise ValueError( + "apply_rope_fusion is not available. Please install TE >= 1.4." + ) + + if self.fused_single_qkv_rope: + if self.attention_output_gate: + raise ValueError("fused_single_qkv_rope does not support gated attention for now.") + + if self.multi_latent_attention and self.rotary_interleaved: + raise ValueError("rotary_interleaved does not work with multi_latent_attention.") + + # Set the embedding init method + if self.embedding_init_method_std is None: + # By default, use the same init std as you use for every other non-output layer. + self.embedding_init_method_std = self.init_method_std + + if self.embedding_init_method is None: + if self.init_method is None or (self.embedding_init_method_std != self.init_method_std): + # In this case, we set both the init method and the embedding init method to + # whatever std value requested (or defaulted) for the embedding_init_layer + self.embedding_init_method = init_method_normal(self.embedding_init_method_std) + else: + # Replicate the current behavior where if you are not changing the std of the + # embedding init differently and the init method is set, we fallback to the + # init method for this layer. Since we are here after an OR we know that + # init_method is not None + self.embedding_init_method = self.init_method + + if self.init_method is None: + self.init_method = init_method_normal(self.init_method_std) + + if self.output_layer_init_method is None: + self.output_layer_init_method = scaled_init_method_normal( + self.init_method_std, + self.num_layers, + multiplier=2.0 if not self.is_hybrid_model else 1.0, + ) + + if self.num_moe_experts is not None and self.add_bias_linear: + assert ( + self.expert_tensor_parallel_size == 1 + ), "Bias in Moe is only supported when ETP==1" + + if self.moe_router_enable_expert_bias and self.moe_router_score_function != "sigmoid": + raise ValueError( + "Expert bias for aux-loss-free routing only supports sigmoid score function." + "Please set --moe-router-score-function sigmoid for sigmoid score function." + ) + + if self.num_moe_experts and self.fp8: + # TE version below 1.7.0 will raise Error when handle zeros tokens for expert + if not is_te_min_version("1.7.0.dev0"): + raise ValueError( + "Only transformer-engine>=1.7.0 supports MoE FP8 training, " + f"but your version is {get_te_version()}." + ) + + if self.moe_grouped_gemm and not is_te_min_version("1.11.0"): + raise ValueError( + "Only transformer-engine>=1.11.0 supports FP8 grouped gemm, " + f"but your version is {get_te_version()}." + ) + + if self.moe_router_padding_for_fp8: + # enable moe_router_padding_for_quantization + warnings.warn( + "--moe-router-padding-for-fp8 is going to be deprecated. " + "Use --moe-router-padding-for-quantization instead." + ) + self.moe_router_padding_for_quantization = True + + if self.moe_router_padding_for_quantization: + if self.fp8 is None and self.fp4 is None: + raise ValueError( + "fp8/fp4 must be specified when moe_router_padding_for_quantization is True." + ) + + if self.moe_token_dispatcher_type in ["allgather", "alltoall_seq"]: + raise ValueError( + "allgather and alltoall_seq dispatcher does not support " + "moe_router_padding_for_quantization." + ) + + if ( + self.moe_router_topk == 1 + and self.moe_router_score_function == "softmax" + and not self.moe_router_pre_softmax + and self.moe_router_load_balancing_type != "sinkhorn" + ): + # Requires applying softmax before selecting the top-k when k is 1, + # since softmax on a [num_tokens, 1] would yield a zero gradient. + raise ValueError("Please use --moe-router-pre-softmax when topk is 1.") + + if self.moe_router_group_topk: + if self.moe_router_topk_limited_devices: + raise ValueError( + "moe_router_topk_limited_devices is deprecated and replaced by " + "moe_router_group_topk and moe_router_num_groups." + ) + if not self.moe_router_num_groups: + raise ValueError( + "When using group limited routing, moe_router_num_groups must be specified." + ) + else: + assert self.num_moe_experts % self.moe_router_num_groups == 0, ( + f"num_moe_experts ({self.num_moe_experts}) should be divisible by " + f"moe_router_num_groups ({self.moe_router_num_groups})." + ) + assert self.moe_router_group_topk <= self.moe_router_num_groups, ( + f"moe_router_group_topk ({self.moe_router_group_topk}) should be smaller than " + f"moe_router_num_groups ({self.moe_router_num_groups})." + ) + elif self.moe_router_topk_limited_devices: + warnings.warn( + "moe_router_topk_limited_devices is deprecated. Use moe_router_group_topk and " + "moe_router_num_groups instead." + ) + self.moe_router_group_topk = self.moe_router_topk_limited_devices + self.moe_router_num_groups = self.expert_model_parallel_size + + if self.enable_cuda_graph or self.external_cuda_graph: + assert ( + self.cuda_graph_impl == "none" + ), "Do not use enable_cuda_graph or external_cuda_graph with cuda_graph_impl." + assert ( + not self.enable_cuda_graph or not self.external_cuda_graph + ), "enable_cuda_graph and external_cuda_graph cannot be enabled at the same time." + + if self.enable_cuda_graph: + warnings.warn('enable_cuda_graph is deprecated, use cuda_graph_impl=local instead.') + self.cuda_graph_impl = "local" + if self.external_cuda_graph: + warnings.warn( + 'external_cuda_graph is deprecated, ' + 'use cuda_graph_impl=transformer_engine instead.' + ) + self.cuda_graph_impl = "transformer_engine" + + if self.cuda_graph_scope is None: + self.cuda_graph_scope = [] + elif not isinstance(self.cuda_graph_scope, list): + if isinstance(self.cuda_graph_scope, CudaGraphScope): + self.cuda_graph_scope = [self.cuda_graph_scope] + else: + assert isinstance(self.cuda_graph_scope, str), ( + "cuda_graph_scope must be a string that can be converted to a list of " + f"CudaGraphScope, got {self.cuda_graph_scope}." + ) + self.cuda_graph_scope = self.cuda_graph_scope.split(',') + if all(isinstance(scope, str) for scope in self.cuda_graph_scope): + # Backward compatibility for "full" scope. Now we use an empty list instead. + if "full" in self.cuda_graph_scope: + assert self.cuda_graph_scope == [ + "full" + ], "full scope cannot be used with other scopes." + warnings.warn( + "full scope is deprecated. " + "Use empty cuda_graph_scope to capture the whole layer." + ) + self.cuda_graph_scope = [] + else: + self.cuda_graph_scope = [CudaGraphScope[scope] for scope in self.cuda_graph_scope] + assert all( + isinstance(scope, CudaGraphScope) for scope in self.cuda_graph_scope + ), f"cuda_graph_scope must be a list of CudaGraphScope, got {self.cuda_graph_scope}." + + if self.cuda_graph_impl != "none": + assert self.cuda_graph_impl in [ + "transformer_engine", + "local", + ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" + + if self.cpu_offloading: + raise ValueError("CUDA graphs not supported with CPU offloading.") + + if self.cuda_graph_impl == "local": + assert not self.cuda_graph_scope or self.cuda_graph_scope == [ + CudaGraphScope.full_iteration + ], ( + "For local cuda graph implementation, the only valid value for " + "cuda_graph_scope is full_iteration, or an empty list to denote layerwise " + "graphs. To use other scopes, use cuda_graph_impl=transformer_engine." + ) + + if self.cuda_graph_impl == "transformer_engine": + assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, ( + "To use full iteration cuda graph, please use " + "cuda_graph_impl=local instead of cuda_graph_impl=transformer_engine." + ) + assert ( + CudaGraphScope.moe not in self.cuda_graph_scope + or CudaGraphScope.moe_router not in self.cuda_graph_scope + ), 'cuda_graph_scope must not contain both moe and moe_router.' + if CudaGraphScope.moe_preprocess in self.cuda_graph_scope: + assert ( + CudaGraphScope.moe_router in self.cuda_graph_scope + ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' + if self.num_moe_experts is None or self.num_moe_experts <= 1: + assert ( + CudaGraphScope.moe not in self.cuda_graph_scope + and CudaGraphScope.moe_router not in self.cuda_graph_scope + ), 'moe cuda graph is only supported for MoE.' + else: + if self.moe_layer_freq == 1 or ( + isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq + ): + assert CudaGraphScope.mlp not in self.cuda_graph_scope, ( + 'mlp cuda graph is only supported for dense layers, ' + 'but not found in the model.' + ) + if ( + self.moe_expert_capacity_factor is None + or not self.moe_pad_expert_input_to_capacity + ): + assert ( + CudaGraphScope.moe not in self.cuda_graph_scope + ), 'moe cuda graph is only supported with drop-padding MoE.' + if self.moe_token_dispatcher_type == 'alltoall' and ( + self.moe_expert_capacity_factor is not None + or self.moe_router_padding_for_quantization + ): + assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, ( + 'moe_preprocess cuda graph is not supported when there are ' + 'DtoH copies and synchronizations in the preprocess step.' + ) + + if self.recompute_granularity: + if self.recompute_granularity != "selective" or not self.cuda_graph_scope: + raise ValueError( + "Full-layer CUDA graphs not supported with activation recomputation." + ) + elif self.cuda_graph_scope != [CudaGraphScope.full_iteration]: + # For scoped CUDA graphs, only the non-graphed parts of the layer can be + # recomputed. So check if there are overlaps between the recomputed parts + # and the graphed parts. + if CudaGraphScope.attn in self.cuda_graph_scope: + for module in self.recompute_modules: + if module in ['core_attn', 'mla_up_proj']: + raise ValueError( + f'attn cuda graph is not supported with {module} recompute.' + ) + if ( + CudaGraphScope.mlp in self.cuda_graph_scope + and "mlp" in self.recompute_modules + ): + raise ValueError(f'mlp cuda graph is not supported with mlp recompute.') + if CudaGraphScope.moe in self.cuda_graph_scope: + for module in self.recompute_modules: + if module in ['moe_act', 'moe', 'shared_experts']: + raise ValueError( + f'moe cuda graph is not supported with {module} recompute.' + ) + if CudaGraphScope.moe_router in self.cuda_graph_scope: + for module in self.recompute_modules: + if module in ['moe', 'shared_experts']: + raise ValueError( + f'moe_router cuda graph is not supported with {module} ' + 'recompute.' + ) + if "layernorm" in self.recompute_modules: + if ( + CudaGraphScope.attn in self.cuda_graph_scope + and CudaGraphScope.mlp in self.cuda_graph_scope + and ( + CudaGraphScope.moe in self.cuda_graph_scope + or CudaGraphScope.moe_router in self.cuda_graph_scope + ) + ): + raise ValueError( + 'cuda graph is not supported with layernorm recompute.' + ) + if CudaGraphScope.attn in self.cuda_graph_scope: + warnings.warn( + "input_layernorm recompute is not supported with attention " + "cudagraph. Will only recompute the pre_mlp_layernorm." + ) + if ( + CudaGraphScope.mlp in self.cuda_graph_scope + or CudaGraphScope.moe in self.cuda_graph_scope + or CudaGraphScope.moe_router in self.cuda_graph_scope + ): + warnings.warn( + "pre_mlp_layernorm recompute is not supported with mlp/moe " + "cudagraph. Will only recompute the input_layernorm." + ) + + if self.moe_token_dispatcher_type in ["allgather"]: + if self.variable_seq_lengths is True: + raise ValueError( + f"Token dispatcher type: {self.moe_token_dispatcher_type} does not support " + f"variable sequence length, please use alltoall dispatcher instead." + ) + + if self.moe_permute_fusion: + from megatron.core.transformer.moe.moe_utils import ( + fused_permute, + fused_permute_with_probs, + fused_sort_chunks_by_index, + fused_sort_chunks_by_index_with_probs, + fused_unpermute, + ) + + if ( + fused_permute is None + or fused_permute_with_probs is None + or fused_sort_chunks_by_index is None + or fused_sort_chunks_by_index_with_probs is None + or fused_unpermute is None + ): + raise ValueError("fused permutation is not available. Please install TE >= 2.1.0.") + + if self.overlap_moe_expert_parallel_comm: + # TODO: remove this after we fix the hang issue with torch version < 2.6.0 + assert is_torch_min_version( + "2.6.0" + ), "A2A Overlap encounters hang issue with torch version < 2.6.0" + if self.pipeline_model_parallel_size > 1: + assert self.virtual_pipeline_model_parallel_size is not None, ( + "If enabling EP A2A overlap, virtual_pipeline_model_parallel_size " + "must be specified when pipeline_model_parallel_size > 1" + ) + # Expert model parallelism requirements + assert ( + self.expert_model_parallel_size > 1 + ), 'overlap_moe_expert_parallel_comm is only supported with expert model parallelism' + assert self.moe_token_dispatcher_type in [ + 'alltoall', + 'flex', + ], 'overlap_moe_expert_parallel_comm is supported with alltoall/flex token dispatcher' + + assert ( + self.recompute_granularity != 'full' + ), 'disable full recomputation when enabling overlap_moe_expert_parallel_comm' + assert ( + self.recompute_method is None + ), 'disable recomputation method when enabling overlap_moe_expert_parallel_comm' + assert ( + self.recompute_num_layers is None + ), 'recompute_num_layers must be None when enabling overlap_moe_expert_parallel_comm' + + # Check if bf16 or fp16 is used + assert ( + self.bf16 or self.fp16 + ), 'overlap_moe_expert_parallel_comm is only supported with bf16 or fp16 model' + + assert ( + not self.moe_shared_expert_overlap + ), 'disable moe_shared_expert_overlap when enabling overlap_moe_expert_parallel_comm' + assert ( + self.mtp_num_layers is None or self.mtp_num_layers == 1 + ), 'MTP layernum only supports 1 when enabling overlap_moe_expert_parallel_comm.' + + # Check delay_wgrad_compute compatibility + if self.delay_wgrad_compute: + assert ( + self.overlap_moe_expert_parallel_comm + ), 'overlap_moe_expert_parallel_comm must be enabled when enabling delay_wgrad_compute' + assert ( + not self.moe_use_legacy_grouped_gemm + ), 'delay_wgrad_compute is not supported with legacy groupedgemm implementation' + + if self.context_parallel_size > 1 and self.cp_comm_type is not None: + if isinstance(self.cp_comm_type, list): + assert len(self.cp_comm_type) == self.num_layers, ( + f"Length of cp_comm_type ({len(self.cp_comm_type)}) should equal to " + f"the total number of transformer layers ({self.num_layers})!" + ) + else: + assert isinstance( + self.cp_comm_type, str + ), "Unsupported communication type for context parallelism!" + + assert ( + self.pipeline_model_parallel_size > 0 + ), f"Pipeline model parallel size must be larger than 0 \ + when enable --standalone-embedding-stage and --standalone-loss-stage" + + if ( + self.num_moe_experts is not None + and self.num_moe_experts >= 32 + and not self.moe_router_dtype + ): + warnings.warn( + "Using a large number of experts (e.g. >=32) without fp32 routing. " + "Consider enabling moe_router_dtype for better numerical stability." + ) + if self.symmetric_ar_type is not None: + if not HAVE_PACKAGING: + raise ImportError( + "packaging is not installed. Please install it with `pip install packaging`." + ) + assert is_torch_min_version("2.7.0a0"), "Must have at least torch version 2.7 or higher" + assert is_te_min_version("2.3.0") or get_te_version() == PkgVersion( + "2.3.0.dev0+39c0e70" + ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce" + + if self.no_rope_freq: + assert not self.flash_decode, "flash_decode cannot be used with no_rope." + if isinstance(self.no_rope_freq, int): + assert self.num_layers % self.no_rope_freq == 0, ( + f"no_rope_freq={self.no_rope_freq} should be " + f"divisible by num_layers={self.num_layers}." + ) + # Convert integer pattern to list pattern + # e.g. no_rope=4 with num_layers=8 becomes [0,0,0,1,0,0,0,1] + pattern = [0] * (self.no_rope_freq - 1) + [1] + self.no_rope_freq = pattern * (self.num_layers // self.no_rope_freq) + else: + assert len(self.no_rope_freq) == self.num_layers, ( + f"Length of no_rope list ({len(self.no_rope_freq)}) must match " + f"the number of layers ({self.num_layers})" + ) + + if self.fallback_to_eager_attn: + assert self.transformer_impl == "transformer_engine", ( + f"fallback_to_eager_attn is only available with transformer_engine implementation," + f" but got {self.transformer_impl=}." + ) + + if self.fallback_to_eager_attn or self.transformer_impl == "local": + if self.context_parallel_size > 1 and self.cp_comm_type is not None: + all_cp_comm_types_are_all_gather = ( + all(item == "all_gather" for item in self.cp_comm_type) + if isinstance(self.cp_comm_type, list) + else self.cp_comm_type == "all_gather" + ) + if not all_cp_comm_types_are_all_gather: + raise ValueError( + f"fallback_to_eager_attn only supports all_gather communication type " + f"for context parallelism, but got {self.cp_comm_type=} instead." + ) + + +@dataclass +class MLATransformerConfig(TransformerConfig): + """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers. + + The initialization function has an argument for each parameter, including those in + ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA. + """ + + multi_latent_attention: bool = True + """Whether to use Multi-Latent Attention.""" + + q_lora_rank: int = 512 + """Rank of Query tensor's low rank representation.""" + + kv_lora_rank: int = 512 + """Rank of Key and Value tensors' low rank representation.""" + + qk_head_dim: int = 128 + """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + + qk_pos_emb_head_dim: int = 64 + """Dimension of the position embedding in the QK projection.""" + + v_head_dim: int = 128 + """Dimension of the head in the V projection.""" + + normalization: str = "RMSNorm" + """Default normalization layer for MLA models is RMSNorm.""" + + rope_type: str = "yarn" + """Type of RoPE to use. Default to yarn, options are rope and yarn.""" + + rotary_base: float = 10000 + """Rotary base for the rotary embeddings, used by rope and yarn.""" + + rotary_percent: float = 1.0 + """Rotary percent for the rotary embeddings, used by rope.""" + + rotary_scaling_factor: float = 40 + """Rotary scaling factor for the rotary embeddings, used by yarn.""" + + original_max_position_embeddings: int = 4096 + """Original maximum position embeddings for the original model, used by yarn.""" + + beta_fast: float = 32 + """Beta fast for YaRN RoPE, used by yarn.""" + + beta_slow: float = 1 + """Beta slow for YaRN RoPE, used by yarn.""" + + mscale: float = 1.0 + """Mscale for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + + mscale_all_dim: float = 0.0 + """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + + cache_mla_latents: bool = False + """Cache the low dimensional tensors for MLA rather than full KV cache. + This is only for the dynamic inference backend and requires that + Flash MLA is installed.""" + + def __post_init__(self): + super().__post_init__() + if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") + + if self.attention_output_gate: + raise NotImplementedError("Output gate is not supported for MLA yet.") + + if self.cache_mla_latents: + assert ( + self.apply_rope_fusion is False + ), "Rope Fusion is not compatible with caching latents" \ No newline at end of file diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index 886b73e90bc..b19c174e96b 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -427,7 +427,7 @@ index 89659a1d7..c69859a04 100644 + return dq, dk, None, None, None, None, None, None \ No newline at end of file diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py -index fc994490b..b23d2e9a8 100644 +index 353b31e9b..f5c55e429 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -6,6 +6,7 @@ from dataclasses import dataclass @@ -438,16 +438,15 @@ index fc994490b..b23d2e9a8 100644 from megatron.core import parallel_state from megatron.core.models.common.embeddings import ( -@@ -21,6 +22,8 @@ from megatron.core.transformer.module import MegatronModule +@@ -20,6 +21,7 @@ from megatron.core.transformer.enums import AttnMaskType + from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import TransformerConfig - +from megatron.core.transformer.dot_product_attention_context_parallel import AllGatherComm, AttentionFuncionWithContextParallel -+ + try: from fast_hadamard_transform import hadamard_transform - except ImportError: -@@ -191,44 +194,72 @@ def compute_dsa_indexer_loss( +@@ -191,44 +193,72 @@ def compute_dsa_indexer_loss( Returns: index_loss: KL divergence loss (scalar). """ @@ -555,9 +554,9 @@ index fc994490b..b23d2e9a8 100644 if pg_collection.tp.size() > 1: # attention scores are scattered to TP ranks in head dimension. torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) -@@ -252,6 +283,57 @@ def compute_dsa_indexer_loss( - return indexer_loss +@@ -251,6 +281,56 @@ def compute_dsa_indexer_loss( + return indexer_loss +def compute_attention_scores_with_cp(q, k, attn_bias, scale, heads_k_stride = 1): + """ @@ -609,11 +608,10 @@ index fc994490b..b23d2e9a8 100644 + + return attns + -+ + class DSAIndexerLossAutoScaler(torch.autograd.Function): """An AutoScaler that triggers the backward pass and scales the grad for indexer loss. - -@@ -496,7 +578,15 @@ class DSAIndexer(MegatronModule): +@@ -496,7 +576,15 @@ class DSAIndexer(MegatronModule): # Compute attention scores: q @ k^T # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T # -> [seqlen_q, batch, index_n_heads, seqlen_k] @@ -630,19 +628,35 @@ index fc994490b..b23d2e9a8 100644 # Apply ReLU activation. index_scores = torch.relu(index_scores) -@@ -606,7 +696,10 @@ class DSAIndexer(MegatronModule): +@@ -546,14 +634,10 @@ class DSAIndexer(MegatronModule): + None, None, x, self.config, packed_seq_params + ) + if self.config.rope_type == "rope": +- rotary_pos_emb = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + mscale = 1.0 + else: +- rotary_pos_emb, mscale = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + + # ========================================= + # Gather inputs if sp is enabled +@@ -610,7 +694,9 @@ class DSAIndexer(MegatronModule): # ========================================= # Select top-k indices # ========================================= - topk_k = min(self.index_topk, seqlen) + cp_size = parallel_state.get_context_parallel_world_size() -+ + seqlen_k_global = k.shape[0] * cp_size + topk_k = min(self.index_topk, seqlen_k_global) # [batch, seqlen, index_topk] topk_indices = index_scores.topk(topk_k, dim=-1)[1] -@@ -687,6 +780,48 @@ def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): +@@ -691,6 +777,48 @@ def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): output = output.reshape(sq, b, np * hnv) return output @@ -691,7 +705,7 @@ index fc994490b..b23d2e9a8 100644 class DSAttention(MegatronModule): """ -@@ -729,7 +864,6 @@ class DSAttention(MegatronModule): +@@ -733,7 +861,6 @@ class DSAttention(MegatronModule): self, query: torch.Tensor, key: torch.Tensor, @@ -699,7 +713,7 @@ index fc994490b..b23d2e9a8 100644 x: torch.Tensor, qr: torch.Tensor, attention_mask: torch.Tensor, -@@ -743,7 +877,6 @@ class DSAttention(MegatronModule): +@@ -747,7 +874,6 @@ class DSAttention(MegatronModule): Args: query: Query tensor [sq, b, np, hn]. key: Key tensor [skv, b, np, hn]. @@ -707,7 +721,7 @@ index fc994490b..b23d2e9a8 100644 x: Original hidden states [sq, b, hidden_size]. qr: Low-rank query representation [sq, b, q_lora_rank]. attention_mask: Attention mask tensor [b, 1, sq, sk]. -@@ -754,9 +887,11 @@ class DSAttention(MegatronModule): +@@ -758,9 +884,11 @@ class DSAttention(MegatronModule): Returns: output: Output tensor [sq, b, hidden_size] """ @@ -722,7 +736,7 @@ index fc994490b..b23d2e9a8 100644 # Detach x and qr to prevent gradients of indexer from flowing back to the main model. x = x.detach() -@@ -768,18 +903,17 @@ class DSAttention(MegatronModule): +@@ -772,18 +900,17 @@ class DSAttention(MegatronModule): # Generate upper triangular mask with -inf above diagonal, 0 elsewhere # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal) # float_mask [sq, skv] @@ -741,7 +755,7 @@ index fc994490b..b23d2e9a8 100644 - float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( - mask, float('-inf') - ) -+ ++ + # float_mask [b, sq, skv] + float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( + mask, float('-inf') @@ -749,7 +763,7 @@ index fc994490b..b23d2e9a8 100644 # =================================== # Get index scores and top-k indices -@@ -791,32 +925,6 @@ class DSAttention(MegatronModule): +@@ -795,32 +922,6 @@ class DSAttention(MegatronModule): # =================================== # Run sparse attention kernel # =================================== @@ -780,30 +794,24 @@ index fc994490b..b23d2e9a8 100644 - ) - # Attach loss to output - output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) +- + output = unfused_dsa_fn_with_cp(query, key, dim_v, topk_indices, self.softmax_scale) - ++ return output diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py -index 3953d933b..7d030ad02 100644 +index ed90fdffa..4298e044b 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py -@@ -6,6 +6,7 @@ from dataclasses import dataclass - from typing import NoReturn, Optional, Union - - import torch -+import torch.nn.functional as F +@@ -15,7 +15,7 @@ except ImportError: + HAVE_EINOPS = False - try: - from einops import rearrange -@@ -167,6 +168,7 @@ class MultiLatentAttention(Attention): - ) - # Output. -+ # SP_Reduce scatter + TP_Row_par - self.linear_proj = build_module( - submodules.linear_proj, - self.query_projection_size, -@@ -311,7 +313,6 @@ class MultiLatentAttention(Attention): +-from megatron.core import tensor_parallel ++from megatron.core import parallel_state, tensor_parallel + from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, +@@ -312,7 +312,6 @@ class MultiLatentAttention(Attention): core_attn_out = self.core_attention( query, key, @@ -811,7 +819,7 @@ index 3953d933b..7d030ad02 100644 x=hidden_states, qr=q_compressed, attention_mask=attention_mask, -@@ -370,6 +371,19 @@ class MultiLatentAttention(Attention): +@@ -371,6 +370,19 @@ class MultiLatentAttention(Attention): self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) self.qkv_up_checkpoint = None @@ -831,27 +839,53 @@ index 3953d933b..7d030ad02 100644 # ================= # Output. [sq, b, h] # ================= -@@ -384,7 +398,6 @@ class MultiLatentAttention(Attention): - - return output, bias - -- - class MLASelfAttention(MultiLatentAttention): - """MLA Self-attention layer class +@@ -555,11 +567,7 @@ class MLASelfAttention(MultiLatentAttention): + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" +- if packed_seq_params is not None: +- assert ( +- packed_seq_params.local_cp_size is None +- ), "hybrid_context_parallel is not supported with MLA yet and is planned for future. \ +- Please disable hybrid_context_parallel." ++ -@@ -753,7 +766,6 @@ class MLASelfAttention(MultiLatentAttention): - # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] - k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + inference_context = deprecate_inference_params(inference_context, inference_params) -- # todo add assert about fusions and caching +@@ -576,13 +584,11 @@ class MLASelfAttention(MultiLatentAttention): + rotary_pos_sin = None + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if self.config.rope_type == "rope": +- rotary_pos_emb = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + else: if self.config.apply_rope_fusion: - cp_rank = self.pg_collection.cp.rank() - cp_size = self.pg_collection.cp.size() -@@ -844,6 +856,98 @@ class MLASelfAttention(MultiLatentAttention): - value = value.contiguous() + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( +- rotary_seq_len, dtype=hidden_states.dtype, packed_seq_params=packed_seq_params ++ rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" +@@ -591,11 +597,9 @@ class MLASelfAttention(MultiLatentAttention): + and fused_apply_mla_rope_for_kv is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: +- rotary_pos_emb, mscale = self.rotary_pos_emb( +- rotary_seq_len, packed_seq_params=packed_seq_params +- ) ++ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + +- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': ++ if packed_seq_params is not None: + if packed_seq_params.cu_seqlens_q_padded is not None: + cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded + else: +@@ -867,6 +871,98 @@ class MLASelfAttention(MultiLatentAttention): return query, key, value -+ + + def mla_absorb(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] @@ -896,7 +930,7 @@ index 3953d933b..7d030ad02 100644 + + # TODO: Does it match ZZ? SP does not need but CP needs + if self.config.sequence_parallel: -+ kv_compressed = gather_from_sequence_parallel_region(kv_compressed) ++ kv_compressed = gather_from_sequence_parallel_region(kv_compressed, group=self.tp_group) + + # kv_compressed: [num_tokens, kv_lora_rank] + if kv_compressed.ndim == 3: # [s, b, kv_lora_rank] @@ -943,10 +977,11 @@ index 3953d933b..7d030ad02 100644 + key = key.contiguous() + + return query, key - ++ if self.recompute_up_proj: quantization = self.config.fp8 or self.config.fp4 -@@ -860,9 +964,10 @@ class MLASelfAttention(MultiLatentAttention): + self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) +@@ -882,9 +978,10 @@ class MLASelfAttention(MultiLatentAttention): q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb ) else: @@ -958,13 +993,11 @@ index 3953d933b..7d030ad02 100644 if return_compressed_tensors: return query, key, value, q_compressed, kv_compressed -@@ -1104,5 +1209,27 @@ class MLASelfAttention(MultiLatentAttention): - * (self.config.qk_head_dim + self.config.v_head_dim), - -1, +@@ -1128,3 +1225,26 @@ class MLASelfAttention(MultiLatentAttention): ) -- + return weight_kv_updated -+ ++ + @property + def up_k_weight_(self): + # linear_kv_up_proj.weight: [num_heads_per_partition * (qk_head_dim + v_head_dim), kv_lora_rank] @@ -987,9 +1020,10 @@ index 3953d933b..7d030ad02 100644 + ) + # [num_heads_per_partition, v_head_dim, kv_lora_rank] + return weight_reshaped[:, self.config.qk_head_dim:, :] +\ No newline at end of file diff --git a/megatron/core/transformer/tilelang_kernel/__init__.py b/megatron/core/transformer/tilelang_kernel/__init__.py new file mode 100644 -index 000000000..d8f2425f0 +index 000000000..c63794256 --- /dev/null +++ b/megatron/core/transformer/tilelang_kernel/__init__.py @@ -0,0 +1,10 @@ @@ -1003,6 +1037,7 @@ index 000000000..d8f2425f0 + "sparse_mla_fwd_interface", + "sparse_mla_bwd", +] +\ No newline at end of file diff --git a/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py b/megatron/core/transformer/tilelang_kernel/sparse_mla_bwd.py new file mode 100644 index 000000000..83a259efa @@ -1481,10 +1516,10 @@ index 000000000..e247038de + return out, lse \ No newline at end of file diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py -index a3a167549..98391fda6 100644 +index e2705bd9f..29a0ff9e0 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py -@@ -918,9 +918,9 @@ class TransformerConfig(ModelParallelConfig): +@@ -935,11 +935,10 @@ class TransformerConfig(ModelParallelConfig): f" but got {self.context_parallel_size=}." ) elif self.experimental_attention_variant == "dsa": @@ -1495,5 +1530,7 @@ index a3a167549..98391fda6 100644 + # self.context_parallel_size == 1 + # ), "Currently context parallelism is not supported by DSAttention!" assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" - +- if self.fp8: + # cannot support first last layer bf16 with delayed scaling + if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed: From f8e4cd884217b47ceec431b5b4f513fcba51bfa5 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Thu, 22 Jan 2026 15:12:17 -0800 Subject: [PATCH 29/30] update --- docker/deepseekv32/megatron.patch | 19 +++++++------------ miles/backends/training_utils/data.py | 7 ------- miles/utils/data.py | 8 +------- 3 files changed, 8 insertions(+), 26 deletions(-) diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index b19c174e96b..bd5ef87c66c 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -427,7 +427,7 @@ index 89659a1d7..c69859a04 100644 + return dq, dk, None, None, None, None, None, None \ No newline at end of file diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py -index 353b31e9b..f5c55e429 100644 +index 353b31e9b..221e93500 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -6,6 +6,7 @@ from dataclasses import dataclass @@ -635,13 +635,13 @@ index 353b31e9b..f5c55e429 100644 - rotary_pos_emb = self.rotary_pos_emb( - rotary_seq_len, packed_seq_params=packed_seq_params - ) -+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) ++ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) mscale = 1.0 else: - rotary_pos_emb, mscale = self.rotary_pos_emb( - rotary_seq_len, packed_seq_params=packed_seq_params - ) -+ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) ++ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) # ========================================= # Gather inputs if sp is enabled @@ -799,7 +799,7 @@ index 353b31e9b..f5c55e429 100644 + return output diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py -index ed90fdffa..4298e044b 100644 +index ed90fdffa..7a7597d66 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -15,7 +15,7 @@ except ImportError: @@ -852,22 +852,17 @@ index ed90fdffa..4298e044b 100644 inference_context = deprecate_inference_params(inference_context, inference_params) -@@ -576,13 +584,11 @@ class MLASelfAttention(MultiLatentAttention): +@@ -576,9 +584,7 @@ class MLASelfAttention(MultiLatentAttention): rotary_pos_sin = None packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' if self.config.rope_type == "rope": - rotary_pos_emb = self.rotary_pos_emb( - rotary_seq_len, packed_seq_params=packed_seq_params - ) -+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) ++ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) else: if self.config.apply_rope_fusion: rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( -- rotary_seq_len, dtype=hidden_states.dtype, packed_seq_params=packed_seq_params -+ rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq - ) - rotary_pos_emb = None - assert inference_context is None, "Inference with MLA RoPE fusion is not supported" @@ -591,11 +597,9 @@ class MLASelfAttention(MultiLatentAttention): and fused_apply_mla_rope_for_kv is not None ), "Fused MLA RoPE apply is not imported successfully" @@ -875,7 +870,7 @@ index ed90fdffa..4298e044b 100644 - rotary_pos_emb, mscale = self.rotary_pos_emb( - rotary_seq_len, packed_seq_params=packed_seq_params - ) -+ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) ++ rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq_params=packed_seq_params) - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if packed_seq_params is not None: diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index d38161e6070..67bb30108d1 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -134,13 +134,6 @@ def get_batch( tokens = [slice_with_cp(t, pad_token_id, parallel_state, qkv_format, max_seqlen) for t in tokens] tokens = torch.stack(tokens) - if qkv_format == "bshd": - max_seqlen = batch["max_seq_lens"][0] - assert max([t.size(0) for t in tokens]) <= max_seqlen - tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens] - tokens = torch.stack(tokens) - # TODO: padding to multiples? - elif qkv_format == "thd": tokens = [slice_with_cp(t, pad_token_id, parallel_state, qkv_format) for t in tokens] diff --git a/miles/utils/data.py b/miles/utils/data.py index 737246acddd..eb512e51481 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -200,13 +200,6 @@ def __init__( metadata["tools"] = tools if apply_chat_template: - output_prompt = tokenizer.apply_chat_template( - prompt, - tools=tools, - tokenize=False, - add_generation_prompt=True, - **(apply_chat_template_kwargs or {}), - ) ### DSV32 try: prompt = tokenizer.apply_chat_template( @@ -221,6 +214,7 @@ def __init__( encode_config = dict(thinking_mode="thinking", drop_thinking=True, add_default_bos_token=True) prompt = encode_messages(prompt, **encode_config) ### DSV32 + output_prompt = prompt else: output_prompt = prompt From 3dfbb7fe931810e1171ba6d92acabd468235fc2d Mon Sep 17 00:00:00 2001 From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com> Date: Sat, 24 Jan 2026 15:46:32 -0800 Subject: [PATCH 30/30] Rebase dsv32 (#516) --- A.py | 1986 ------------------------ docker/deepseekv32/megatron.patch | 28 + miles/backends/megatron_utils/actor.py | 2 +- miles/utils/data.py | 3 +- scripts/run_deepseek_v32.py | 15 +- 5 files changed, 40 insertions(+), 1994 deletions(-) delete mode 100644 A.py diff --git a/A.py b/A.py deleted file mode 100644 index 811f1c692d6..00000000000 --- a/A.py +++ /dev/null @@ -1,1986 +0,0 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -import warnings -from dataclasses import dataclass -from typing import Callable, List, Literal, Optional, Tuple, Union - -import torch -import torch.nn.functional as F - -from megatron.core.enums import Fp4Recipe, Fp8Recipe -from megatron.core.quantization.quant_config import RecipeConfig -from megatron.core.transformer.enums import AttnBackend, CudaGraphScope -from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout - -from ..fusions.fused_bias_geglu import quick_gelu -from ..model_parallel_config import ModelParallelConfig -from ..utils import ( - get_te_version, - init_method_normal, - is_te_min_version, - is_torch_min_version, - scaled_init_method_normal, -) - -try: - from packaging.version import Version as PkgVersion - - HAVE_PACKAGING = True -except ImportError: - HAVE_PACKAGING = False - - -@dataclass -class TransformerConfig(ModelParallelConfig): - """Configuration object for megatron-core transformers. - - The initialization function has an argument for each parameter, - including those in ModelParallelConfig. - """ - - #################### - # model architecture - #################### - - num_layers: int = 0 - """Number of transformer layers in a transformer block.""" - - mtp_num_layers: Optional[int] = None - """Number of Multi-Token Prediction (MTP) Layers.""" - - mtp_loss_scaling_factor: Optional[float] = None - """Weighting factor of Multi-Token Prediction (MTP) loss.""" - - num_layers_in_first_pipeline_stage: Optional[int] = None - """Number of transformer layers on first pipeline stage. - None implies equal layer division across PP ranks.""" - - num_layers_in_last_pipeline_stage: Optional[int] = None - """Number of transformer layers on last pipeline stage. - None implies equal layer division across PP ranks.""" - - pipeline_model_parallel_layout: Optional[Union[str, list, PipelineParallelLayerLayout]] = None - """Custom definition of the pipeline parallel partitioning. - Support type: - - str: e.g., 'Et*3|(tt|)*29,m|L'. Stages are split by '|', replicated stages or layers - can be described with multiplication. Commas can be used cosmetically. - - list: e.g., [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]. - - PipelineParallelLayerLayout: a PipelineParallelLayerLayout object. - If given either a string or a list, it will be transferred into a PipelineParallelLayerLayout - in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers - in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., - vpp(i)pp(j), vpp(i)pp(j+1), ..., vpp(-1)pp(-2), vpp(-1)pp(-1). - In the inner lists of layers, 'embedding' or 'E' denotes the embedding layer, 'loss' or 'L' - denotes the loss function, and 'decoder' or 't' denotes the transformer decoder layer. - Examples: - [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]: - pp = 2, vpp = None - pp rank 0 holds: embedding, decoder - pp rank 1 holds: decoder*3, loss - 'E|(tt|)*2,(t|)*4,mL': - pp = 2, vpp = 4 - vpp rank 0 pp rank 0 holds: embedding - vpp rank 0 pp rank 1~2 holds: decoder*2 - vpp rank 0 pp rank 3 holds: decoder - vpp rank 1 pp rank 0~2 holds: decoder - vpp rank 1 pp rank 3 holds: mtp, loss""" - - account_for_embedding_in_pipeline_split: bool = False - """If set, the embedding layer will be treated as a standard transformer - layer in the context of partition and placement for pipeline parallelism.""" - - account_for_loss_in_pipeline_split: bool = False - """If set, the loss layer will be treated as a standard transformer - layer in the context of partition and placement for pipeline parallelism.""" - - hidden_size: int = 0 - """Transformer hidden size.""" - - num_attention_heads: int = 0 - """Number of transformer attention heads.""" - - attention_backend: AttnBackend = AttnBackend.auto - """Attention backend to run. By default we let transformer engine - decide the best backend to run (except in the case of local). - If attention backend is local we use the local pytorch implementation in mcore. - Users can specify exact backend by changing this config. """ - - softmax_scale: Optional[float] = None - """Softmax scale for attention scaling.""" - - softmax_type: Literal['vanilla', 'off-by-one', 'learnable'] = 'vanilla' - """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. - Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and - and learnable offset.""" - - num_query_groups: Optional[int] = None - """Number of query groups for group query attention. If None, normal attention is used.""" - - ffn_hidden_size: Optional[int] = None - """Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size - if not provided.""" - - kv_channels: Optional[int] = None - """Projection weights dimension in multi-head attention. This is set to hidden_size // - num_attention_heads if not provided.""" - - hidden_dropout: float = 0.1 - """Dropout probability for transformer hidden state.""" - - attention_dropout: float = 0.1 - """Post attention dropout probability.""" - - fp32_residual_connection: bool = False - """If true, move residual connections to fp32.""" - - # @jcasper should we keep this option? - apply_residual_connection_post_layernorm: bool = False - """If True, uses the original BERT residule connection ordering.""" - - layernorm_epsilon: float = 1e-5 - """Epsilon value for any LayerNorm operations.""" - - layernorm_zero_centered_gamma: bool = False - """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves - numerical stability.""" - - add_bias_linear: bool = True - """Include a bias term in all linear layers (QKV projections, after core attention, and two in - MLP layer).""" - - add_qkv_bias: bool = False - """Add a bias term only for QKV projections.""" - - gated_linear_unit: bool = False - """Use a gated linear unit for the first linear layer in the MLP.""" - - activation_func: Callable = F.gelu - """Activation function to use for the non-linearity in the MLP.""" - - activation_func_fp8_input_store: bool = False - """Store the input of MLP activation function in FP8 for backprop to save memory. - The stored input is casted back to the original precision before backprop compuatation.""" - - glu_linear_offset: float = 0.0 - """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only - used when gated_linear_unit is True""" - - activation_func_clamp_value: Optional[float] = None - """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func - is quick_gelu.""" - - num_moe_experts: Optional[int] = None - """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None - for no MoE.""" - - rotary_interleaved: bool = False - """True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of - first half and second half (LLaMa style). Default to False.""" - - window_size: Optional[Tuple[int, int]] = None - """If not None, then will use sliding window attention. The size of the window is specified by - the numbers inside the tuple; -1 is special value meaning "infinite window size".""" - - window_attn_skip_freq: Optional[Union[int, List[int]]] = None - """Frequency of full attention layers among sliding window attention layers. Accepts either: - - An integer N: Represents a (N-1):1 ratio, one full attention layer after (N-1) SWA layers. - - A list that defines a custom pattern, e.g.: [1,1,1,1,0,0,0,0], where 1 represents SWA. """ - - normalization: str = "LayerNorm" - """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`.""" - - qk_layernorm: bool = False - """Whether to apply `normalization` type of normalization to the query and key embeddings.""" - - qk_clip: bool = False - """Whether to clip the query and key weights. Needed for Muon MLA Model training.""" - - qk_clip_alpha: float = 0.5 - """The balancing alpha for qk-clip. Q = Q * (eta ** alpha)""" - - qk_clip_threshold: float = 100 - """The balancing threshold for qk-clip. eta = min(threshold / max_attention_logits, 1.0)""" - - log_max_attention_logit: bool = False - """Whether to log the max attention logit across whole model. Decoupled from qk_clip, - defualts to False. Setting qk_clip will automatically log the max logit""" - - attention_output_gate: bool = False - """Whether to apply output gate to the attention layers.""" - - test_mode: bool = False - """Whether to run real-time tests.""" - - calculate_per_token_loss: bool = False - """Whether cross entropy loss is calculated over the actual number of non-padded tokens in the - global batch, versus the default behavior of assuming all tokens are non-padded.""" - - multi_latent_attention: bool = False - """Whether to use multi-latent attention.""" - - no_rope_freq: Optional[Union[int, List[int]]] = None - """Controls which layers perform Rotary Position Embedding (RoPE). Accepts either: - An integer N: Creates a pattern where RoPE is skipped every N-1 layers. For example, - no_rope=4 means RoPE is applied for 3 layers, then skipped for 1 layer, repeating this pattern. - A list of integers: Defines a custom pattern where 1 means skip RoPE and 0 means apply RoPE. - For example, [0,1,1,0] means: apply RoPE, skip RoPE, skip RoPE, apply RoPE.""" - - moe_deepep_num_sms: int = 20 - """Number of SMs to use for DeepEP.""" - - moe_hybridep_num_sms: int = 16 - """Number of SMs to use for HybridEP. In pure NVL scenarios, - 16 SMs can generally achieve good bandwidth.""" - - #################### - # attention variant - #################### - experimental_attention_variant: Optional[str] = None - """Type of attention variant to use. Currently support gated_delta_net and dsa.""" - - #################### - # attention variant: gated_delta_net - #################### - linear_attention_freq: Optional[Union[int, List[int]]] = None - """Frequency between LA (linear attention) layers - and SDPA (scaled dot-product attention) layers. - Accepts either: - - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer - - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]""" - - linear_conv_kernel_dim: Optional[int] = None - """Conv kernel dimension for the gated delta net.""" - - linear_key_head_dim: Optional[int] = None - """Query and key head dimension for the gated delta net.""" - - linear_value_head_dim: Optional[int] = None - """Value and gate head dimension for the gated delta net.""" - - linear_num_key_heads: Optional[int] = None - """Number of query and key heads for the gated delta net.""" - - linear_num_value_heads: Optional[int] = None - """Number of value and gate heads for the gated delta net.""" - - #################### - # attention variant: dsa - #################### - dsa_indexer_n_heads: Optional[int] = None - """Number of DSA indexer heads.""" - - dsa_indexer_head_dim: Optional[int] = None - """Dimension per DSA indexer head.""" - - dsa_indexer_topk: Optional[int] = None - """Number of top-k tokens to select in DSA indexer.""" - - dsa_indexer_loss_coeff: Optional[float] = None - """Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss.""" - - dsa_indexer_use_sparse_loss: Optional[bool] = None - """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the - top-k indices.""" - - #################### - # initialization - #################### - init_method: Optional[Callable] = None - """Method to initialize weights. Note that bias is always set to zero. Should be a function that - takes a single Tensor and initializes it. If None, will be set to - megatron.core.utils.init_method_normal(init_method_std) which is torch nn init normal with - mean=0.0 and std=init_method_std.""" - - output_layer_init_method: Optional[Callable] = None - """Method to initialize weights of the output layer of both attention and MLP blocks. If None, - will be set to megatron.core.utils.scaled_init_method_normal(init_method_std) which is torch nn - init normal with mean=0.0 and std=init_method_std / math.sqrt(2.0 * num_layers).""" - - init_method_std: float = 0.02 - """Standard deviation of the zero mean normal for the default initialization method, not used if - init_method and output_layer_init_method are provided.""" - - embedding_init_method: Optional[Callable] = None - """ - Method to initialize weights of the embedding layer. If None, will be set as described - in init_method above. - """ - - embedding_init_method_std: Optional[float] = None - """ - Standard deviation of the zero mean normal for the default initialization method for the - embedding layer. If None, will be set to init_method_std. - """ - - init_model_with_meta_device: bool = False - """ - If True, initializes the model with the meta device. This is helpful for - training of very large models. This feature is only works when megatron fsdp is turned on. - """ - - #################### - # mixed-precision - #################### - apply_query_key_layer_scaling: bool = False - """If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with - fp16.""" - - attention_softmax_in_fp32: bool = True - """If True, run attention masking and softmax in fp32. This should be True if - apply_query_key_layer_scaling is True.""" - - disable_bf16_reduced_precision_matmul: bool = False - """If True, sets torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=False to - prevent matmul from using reduced precision accumulation when using BF16.""" - - #################### - # fusion - #################### - bias_activation_fusion: bool = False - """If True, fuses bias addition and the activation function when possible.""" - - masked_softmax_fusion: bool = False - """If True, uses softmax fusion.""" - - persist_layer_norm: bool = False - """If True, uses the persistent fused layer norm kernel. This kernel only supports a fixed set - of hidden sizes.""" - - memory_efficient_layer_norm: bool = False - """If True, and using local layers (not from TransformerEngine), tells Apex to use the memory - efficient fused LayerNorm kernel. Ignored if not using LayerNorm.""" - - bias_dropout_fusion: bool = False # TODO: this should be bias_dropout_add_fusion? - """If True, uses bias dropout fusion.""" - - apply_rope_fusion: bool = False - """If True, use fused RoPE kernel.""" - - use_fused_weighted_squared_relu: bool = False - """If True, uses fused weighted squared relu kernel when using MoE.""" - - fused_single_qkv_rope: bool = False - """If set, avoid splitting QKV before ROPE forward and avoid concatenating ROPE dgrads.""" - - #################### - # activation recomputation - #################### - recompute_granularity: Optional[str] = None - """Determines which type of activation recompute to use. Megatron-core supports 'selective' - activation checkpointing where the submodules set in --recompute-modules is checkpointed. - The default is "core_attn" which is the memory intensive part of attention. - These memory intensive activations are also less compute intensive which makes activation - checkpointing more efficient for LLMs (20B+). See Reducing Activation Recomputation in Large - Transformer Models (https://arxiv.org/abs/2205.05198) for more details. 'full' will checkpoint - the entire transformer layer. If None, no recompute is performed and all activations are saved. - If set, must be 'selective' or 'full'. 'selective' always uses all layers. - """ - - recompute_method: Optional[str] = None - """Determines which transformer layers will be recomputed. uniform will uniformly divide the - total number of transformer layers in a transformer block and recompute the input activation of - each divided chunk at the specified granularity. block will recompute the input activations for - only a set number of transformer layers per pipeline stage. The rest of the layers in the - pipeline stage will not have any activations recomputed. If None, and recompute is enabled, all - layers will do recomputation. If set, must be 'uniform' or 'block'.""" - - recompute_num_layers: Optional[int] = None - """When recompute_method is uniform, recompute_num_layers is the number of transformer layers in - each uniformly divided recompute unit. When recompute_method is block, recompute_num_layers is - the number of transformer layers to recompute within each pipeline stage. Must be None for - 'selective' activation checkpointing.""" - - distribute_saved_activations: Optional[bool] = None - """If True, distribute recomputed activations across the model parallel group.""" - - recompute_modules: Optional[List[str]] = None - """The submodules to recompute. - choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts". - default: ["core_attn"]. - "core_attn": recompute the core attention part of the transformer layer. - "moe_act": recompute the MoE MLP activation function. - "layernorm": recompute the input_layernorm and pre_mlp_layernorm. - "mla_up_proj": recompute the MLA up projection and RoPE applying parts. - "mlp": recompute the dense MLP submodule. - "moe": recompute the MoE layer. - "shared_experts": recompute the shared experts in the MoE layer. - "moe_act", "layernorm", and "mla_up_proj" use output-discarding checkpointing, - "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. - """ - - #################### - # fp8 related - #################### - fp8: Optional[str] = None - """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined - choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8 - activation and weight tensors and e5m2 for all FP8 output activation gradient tensors.""" - - fp8_recipe: Optional[str] = "delayed" - """If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined - choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed' - uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only, - 4) 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe.""" - - fp8_param: bool = False - """If set, keep the parameters in fp8 precision to save memory. This option must be used - together with fp8 mode (i.e., TransformerConfig.fp8 is not None). Note that not all parameters - will be converted to fp8; for example, biases will remain unchanged. The parameters affected are - primarily the weights of GEMMs. The specific parameters that will be converted to fp8 are - determined by TE.""" - - fp8_quantizer_factory: Optional[str] = None - """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. - Required when fp8_recipe is custom.""" - - fp8_margin: int = 0 - """Margin for the scaling factor computation.""" - - fp8_interval: int = 1 - """DEPRECATED from TransformerEngine v1.8.0. This flag is ignored. - Controls how often the scaling factor is recomputed. - """ - - fp8_amax_history_len: int = 1 - """The length of the amax history window used for scaling factor computation.""" - - fp8_amax_compute_algo: str = "most_recent" - """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2 - predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent` - always chooses the most recently seen value. - - """ - - fp8_wgrad: bool = True - """When set to False, override FP8 config options and do the wgrad computation - in higher precision.""" - - fp8_dot_product_attention: bool = False - """When set to True, use the FP8 implementation of Dot Product Attention.""" - - fp8_multi_head_attention: bool = False - """When set to True, use the FP8 implementation of Multi Head Attention.""" - - tp_only_amax_red: bool = False - """When set to True, reduce the FP8 AMAX only in the TP or TP-CP domain""" - - first_last_layers_bf16: bool = False - """If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8.""" - - num_layers_at_start_in_bf16: int = 1 - """Number of layers at the start of the model to keep in BF16 precision when - first_last_layers_bf16 is True.""" - - num_layers_at_end_in_bf16: int = 1 - """Number of layers at the end of the model to keep in BF16 precision when - first_last_layers_bf16 is True.""" - - use_kitchen: bool = False - """Use the kitchen extension for transformer quantization.""" - - #################### - # fp4 related - #################### - fp4: Optional[str] = None - """If set, enables the use of FP4 precision through Transformer Engine. Currently only - supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0).""" - - fp4_recipe: Optional[str] = "nvfp4" - """If set, enables the use of FP4 precision through Transformer Engine. Currently only - 'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture.""" - - fp4_param: bool = False - """If set, keep the parameters in fp4 precision to save memory. This option must be used - together with fp4 mode (i.e., TransformerConfig.fp4 is not None). Note that not all parameters - will be converted to fp4; for example, biases will remain unchanged.""" - - fp4_quantizer_factory: Optional[str] = None - """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. - Required when fp4_recipe is custom.""" - - #################### - # MoE related - #################### - moe_shared_expert_intermediate_size: Optional[int] = None - """Shared expert total ffn hidden size. - It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if - there are multiple shared experts. - None means no shared expert. - By default, the shared experts execute before the router. However, when - moe_shared_expert_overlap or overlap_moe_expert_parallel_comm is set, - the shared experts execute after the router, before the routed experts. - This makes the gradients from the router and the shared experts added in - different orders to the hidden_states, causing minor numerical differences - in the hidden_states gradient.""" - - moe_shared_expert_gate: bool = False - """Enable gate for shared expert.""" - - moe_shared_expert_overlap: bool = False - """Enable overlapping between shared expert computations and dispatcher communications. - Without this, the shared experts execute before the router.""" - - moe_layer_freq: Union[int, List[int]] = 1 - """Frequency between MoE layers and Dense layers. Accepts either: - - An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers. - - A list that defines a custom pattern, e.g.: [1,1,1,0,1,1,1,0,1,1,1,0]""" - - moe_ffn_hidden_size: Optional[int] = None - """MoE Feed-Forward Network hidden size""" - - moe_router_load_balancing_type: Union[str, List[str]] = "aux_loss" - """The load balancing strategy for the router. - Options: - - "aux_loss": Load balancing loss used in GShard and SwitchTransformer, calculated at - micro-batch level. - - "seq_aux_loss": Load balancing loss used in DeepSeekV2 and DeepSeekV3, computes loss - for each individual sample. - - "global_aux_loss": Load balancing loss calculated at global batch level. - - "sinkhorn": Balancing algorithm used in S-BASE. - - "none": No load balancing. - A list of strings can be provided to combine multiple aux-loss load balancing types. - The default is "aux_loss". - """ - - moe_router_topk: int = 2 - """Number of experts to route to for each token.""" - - moe_router_topk_limited_devices: Optional[int] = None - """Number of EP ranks to consider for each token in group-limited routing, - DEPRECATED and replaced by moe_router_num_groups and moe_router_group_topk. - """ - - moe_router_padding_for_quantization: Optional[bool] = False - """Whether to pad the routing_map to make sure the number of tokens each expert receives - is a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the explicit - padding in the GroupedMLP layer.""" - - moe_router_padding_for_fp8: Optional[bool] = False - """[Compatibility alias for moe_router_padding_for_quantization] - Enabling this will also enable moe_router_padding_for_quantization.""" - - moe_router_num_groups: Optional[int] = None - """Number of groups to divide experts into for group-limited routing. - When using group-limited routing: - 1. Experts are divided into 'moe_router_num_groups' equal-sized groups - 2. For each token, 'moe_router_group_topk' groups are selected based on sum of - top-('moe_router_topk'/'moe_router_group_topk') routing scores within each group - 3. From these selected groups, 'moe_router_topk' individual experts are chosen - Two common use cases: - - Device-limited routing: Set 'moe_router_num_groups' equal to expert parallel size (EP) - to limit each token to experts on a subset of devices - (See DeepSeek-V2: https://arxiv.org/pdf/2405.04434) - - Node-limited routing: Set 'moe_router_num_groups' equal to number of nodes in EP group - to limit each token to experts on a subset of nodes - (See DeepSeek-V3: https://arxiv.org/pdf/2412.19437) - """ - - moe_router_group_topk: Optional[int] = None - """Number of selected groups for group-limited routing.""" - - moe_router_pre_softmax: bool = False - """Enable pre-softmax(pre-sigmoid) routing for MoE, which means softmax is before the - top-k selection. - By default, softmax is done after top-k.""" - - moe_router_topk_scaling_factor: Optional[float] = None - """Scaling factor for routing score in top-k selection, only works when moe_router_pre_softmax - enabled. Defaults to None, which means no scaling.""" - - moe_router_score_function: str = "softmax" - """Score function for MoE routing. Can be "softmax" or "sigmoid".""" - - moe_router_dtype: Optional[str] = None - """Data type for routing and expert output weighted averaging. Using fp32 or fp64 can - improve stability especially when the number of experts is large (e.g. finegrained-moe). - None means no changes for dtype.""" - - moe_router_enable_expert_bias: bool = False - """TopK routing with dynamic per-expert bias in the aux-loss-free load balancing strategy. - The routing decision is based on the sum of the routing scores and the expert bias. - See https://arxiv.org/abs/2408.15664 for details.""" - - moe_router_bias_update_rate: float = 1e-3 - """The expert bias is updated based on the number of assigned tokens to each expert - in a global batch, where the bias is increased for the experts with less assigned tokens - and decreased for the experts with more assigned tokens. - The default value 1e-3 is same as that used in DeepSeekV3.""" - - moe_router_force_load_balancing: bool = False - """[Experimental] Force load balancing with random logits for MoE router, supports naive topk - and group-limited topk. This is an experimental feature and only for benchmark.""" - - moe_grouped_gemm: bool = False - """When there are multiple experts per rank, compress multiple local (potentially small) gemms - in a single kernel launch to improve the utilization and performance by leveraging the Grouped - GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). - """ - - moe_use_legacy_grouped_gemm: bool = False - """Use legacy GroupedMLP rather than TEGroupedMLP. - Note: The legacy one will be deprecated soon.""" - - moe_aux_loss_coeff: Union[float, List[float]] = 0.0 - """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. - If a list of load balancing types is provided for `moe_router_load_balancing_type`, - a corresponding list of coefficients should be provided here.""" - - moe_z_loss_coeff: Optional[float] = None # 1e-3 would be a good start value for z-loss - """Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended.""" - - moe_input_jitter_eps: Optional[float] = None - """Add noise to the input tensor by applying jitter with a specified epsilon value.""" - - moe_token_dropping: bool = False - """This feature involves selectively dropping and padding tokens for each expert to achieve a - specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is - currently unsupported so should remain False.""" - - moe_token_dispatcher_type: str = "allgather" - """The type of token dispatcher to use. The default is 'allgather'. - Options are 'allgather','alltoall' and 'flex'.""" - - moe_enable_deepep: bool = False - """[Experimental] Enable DeepEP for efficient token dispatching and combine in MoE models.""" - - moe_flex_dispatcher_backend: str = "deepep" - """[Experimental] The backend to use for flex token dispatcher. The default is "deepep". - Options are "deepep" and "hybridep". Currently only "hybridep" backend supports - the MNNVL case.""" - - moe_per_layer_logging: bool = False - """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.""" - - moe_expert_capacity_factor: Optional[float] = None - """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token - will be dropped. The default is None.""" - - moe_pad_expert_input_to_capacity: bool = False - """moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match - the expert capacity length, effective only after the moe_expert_capacity_factor is set. The - default setting is False.""" - - moe_token_drop_policy: str = "probs" - """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with - the lowest probabilities will be dropped. If "position", tokens at the end of each batch will - be dropped. - """ - - moe_layer_recompute: bool = False - """Memory optimization: checkpointing moe_layer to save actiavtion memory.""" - - moe_permute_fusion: bool = False - """Fuse token rearrangement ops during token dispatching.""" - - moe_router_fusion: bool = False - """Fuse ops in routing and aux loss calculation.""" - - moe_apply_probs_on_input: bool = False - """Apply probs on input of experts instead of applying after activation and glu.""" - - ################## - # Context Parallel - ################## - cp_comm_type: Optional[Union[str, List[str]]] = None - """Inter-gpu communication type for context parallelism. - str: all layers share same communication type. - List[str]: each layer has its separate communication type. - cp_comm_type of each layer can be "p2p" or "all_gather" or "a2a" or "a2a+p2p". - "p2p": Exchange KV chunks with P2P communications in ring topology. P2P is async and can be - overlapped with attention compute. - "all_gather": All-gather to get full sequence of KV before attention. The all-gather is not - async, and cannot be overlapped. - "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP group, and gather to get - full sequence of QKV. - "a2a+p2p": A hierarchical implementation of context parallelism to attention. - It uses A2A communications in low-level CP groups (e.g., via NVLink), - and P2P communications in high-level CP groups (e.g., via IBLink). - """ - - ################## - # Cuda Graphs - ################## - enable_cuda_graph: bool = False - """DEPRECATED and replaced by cuda_graph_impl. - When set to true, either partial CUDA graph (1/many CUDA graph per layer) or full iteration - CUDA graph (1 CUDA graph for whole iteration excluding optimizer) is enabled. --cuda-graph-scope - determines the scope of graph capture.""" - - cuda_graph_use_single_mempool: bool = False - """When set to true, cudagraphs will be captured inside a single mempool, in which all - cudagraphs may only be used once per step. If false, cudagraphs may be reused across - microbatches. Enabling may reduce cudagraph memory overheads due to memory fragmentation, - however may greatly increase the number of cudagraphs created when the number of microbatches - is high.""" - - cuda_graph_retain_backward_graph: bool = False - """When set to true, cudagraph backward passes will be graph captured with 'retain_grad=True' - This may enable cudagraphs for certain modules that are not completely cudagraph safe. For - more details, see: https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html.""" - - cuda_graph_warmup_steps: int = 3 - """Number of warmup steps for CUDA graphs""" - - external_cuda_graph: bool = False - """DEPRECATED and replaced by cuda_graph_impl. - When set to true, TransformerLayer layers are swapped with user provided CUDA graphs.""" - - cuda_graph_impl: str = "none" - """Determines the CUDA graph capture implementation. - "none": no CUDA graph. - "local": capture the CUDA graph using MCore local implementation. Either partial CUDA graph - (1/many CUDA graph per layer) or full iteration CUDA graph (1 CUDA graph for whole iteration - excluding optimizer) is enabled. - "transformer_engine": capture the CUDA graph using TE make_graphed_callables().""" - - cuda_graph_scope: Optional[List[CudaGraphScope]] = None - """Determines the CUDA graphs capturing scope. - When cuda_graph_impl is set to "transformer_engine", valid values are "attn", "mlp", "moe", - "moe_router", "moe_preprocess", "mamba". None means the full layer. - When cuda_graph_impl is set to "local", "full_iteration" can be specified as cuda_graph_scope - to enable whole iteration CUDA graph. All other values enable layerwise CUDA graph.""" - - #################### - # miscellaneous - #################### - clone_scatter_output_in_embedding: bool = True - """When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer - to facilitate garbage collection of input.""" - - disable_parameter_transpose_cache: bool = False - """When set to true, the parameter transposes are not cached for subsequent iterations.""" - - config_logger_dir: str = "" - """When non-empty, dumps entry-point configs to config_logger_dir""" - - flash_decode: bool = False - """ Use the optimized flash decoding kernel during inference. """ - - use_te_activation_func: bool = False - """Whether to use ffn activation functions implemented by TransformerEngine""" - - use_te_rng_tracker: bool = False - """ Whether to use the TE or MCore version of the RNG tracker. """ - - inference_rng_tracker: bool = False - """ Whether we should instantiate a separate RNG tracker for inference. """ - - inference_sampling_seed: int = 42 - """ Random seed to use for sampling during inference. """ - - symmetric_ar_type: Optional[str] = None - """Type of symmetric all reduce to use""" - - mrope_section: Optional[List[int]] = None - """ Multimodal rope section is for channel dimension of temporal, height and width - in rope calculation. """ - - is_hybrid_model: bool = False - """ Indicates whether this is a hybrid model. """ - - mamba_state_dim: int = 128 - """The dimensionality of the state representation in Mamba layers.""" - - mamba_head_dim: int = 64 - """The dimensionality of the heads in the Mamba layers.""" - - mamba_num_groups: int = 8 - """The number of groups used in Mamba layers.""" - - mamba_num_heads: Optional[int] = None - """The number of heads used in Mamba layers. - If None, the number of heads will be hidden_size * expand // mamba_head_dim.""" - - use_mamba_mem_eff_path: bool = True - """If True, use the memory efficient path for Mamba layers.""" - - mlp_chunks_for_prefill: int = 1 - """The number of chunks along the sequence dimension to use for MLP computation - during prefill.""" - - heterogeneous_block_specs: bool = False - """Whether to use heterogeneous block specs (nemotron-nas architecture).""" - - hetereogenous_dist_checkpoint: bool = False - """Whether to use heterogenous layers in distributed checkpoint.""" - - #################### - # Quantization - #################### - quant_recipe: Optional[RecipeConfig] = None - """Configuration of any quantization to be applied to the model""" - - transformer_impl: str = "transformer_engine" - """Transformer implementation to use. - Options are 'transformer_engine' for Transformer Engine and 'local' for MCore.""" - - fallback_to_eager_attn: bool = False - """Whether to fallback to eager attention in TE implementation. - Suggested for when desired features are not available in TE implementation.""" - - ##################################### - # Fine-grained Activation Offloading - ##################################### - fine_grained_activation_offloading: bool = False - """If True, offload the input of the specified modules to the CPU. - Fine-grained activation offloading is a module-level offloading method - instead of a layer-level offloading method like cpu_offloading.""" - - offload_modules: Optional[list[str]] = None - """The submodules to offload its input. - choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", - "mlp_norm", "expert_fc1", "moe_act". - "attn_norm": offload the input of the normalization in the attention part. - "qkv_linear": offload the input of the qkv linear part. - "core_attn": offload the input of the core attention part. - "attn_proj": offload the input of the attn linear projection part. - "mlp_norm": offload the input of the normalization in the mlp part. - "expert_fc1": offload the input of the expert fc1 part. - "moe_act": offload the input of the moe act part. - """ - min_offloaded_tensor_size: int = 1024 * 1024 - """The minimum size of the tensor to be offloaded.""" - - def __post_init__(self): - """Python dataclass method that is used to modify attributes after initialization. - See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more - details. - """ - super().__post_init__() - if self.fp16 and self.bf16: - raise ValueError( - f"Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True." - ) - - # Apply BF16 matmul precision setting if needed - if self.bf16 and self.disable_bf16_reduced_precision_matmul: - torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False - - if self.num_attention_heads % self.tensor_model_parallel_size != 0: - raise ValueError( - f"num_attention_heads ({self.num_attention_heads}) must be a multiple of " - f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." - ) - - if self.ffn_hidden_size is None: - self.ffn_hidden_size = 4 * self.hidden_size - - if self.kv_channels is None: - self.kv_channels = self.hidden_size // self.num_attention_heads - - if self.num_query_groups is None: - self.num_query_groups = self.num_attention_heads - - if self.num_query_groups % self.tensor_model_parallel_size != 0: - raise ValueError( - f"num_query_groups ({self.num_query_groups}) must be a multiple of " - f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." - ) - - if self.experimental_attention_variant in ["gated_delta_net"]: - assert ( - self.linear_attention_freq is not None - ), f"linear_attention_freq must be set for linear attention." - - if self.experimental_attention_variant == "gated_delta_net": - # Check required parameters - assert ( - self.linear_conv_kernel_dim is not None - ), "linear_conv_kernel_dim must be set for gated delta net." - assert ( - self.linear_key_head_dim is not None - ), "linear_key_head_dim must be set for gated delta net." - assert ( - self.linear_value_head_dim is not None - ), "linear_value_head_dim must be set for gated delta net." - assert ( - self.linear_num_key_heads is not None - ), "linear_num_key_heads must be set for gated delta net." - assert ( - self.linear_num_value_heads is not None - ), "linear_num_value_heads must be set for gated delta net." - assert self.linear_num_value_heads % self.linear_num_key_heads == 0, ( - f"linear_num_value_heads ({self.linear_num_value_heads}) must be a multiple of " - f"linear_num_key_heads ({self.linear_num_key_heads})." - ) - - # Check tensor parallelism compatibility - assert ( - self.linear_num_key_heads % self.tensor_model_parallel_size == 0 - ), "linear_num_key_heads must be a multiple of tensor_model_parallel_size." - assert ( - self.linear_num_value_heads % self.tensor_model_parallel_size == 0 - ), "linear_num_value_heads must be a multiple of tensor_model_parallel_size." - - # Do not support yet, but coming soon. - assert self.context_parallel_size == 1, ( - f"Gated delta net does not support context parallel for now," - f" but got {self.context_parallel_size=}." - ) - elif self.experimental_attention_variant == "dsa": - # assert ( - # self.context_parallel_size == 1 - # ), "Currently context parallelism is not supported by DSAttention!" - assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" - - if self.fp8: - # cannot support first last layer bf16 with delayed scaling - if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed: - raise ValueError("Delayed scaling does not support first / last layer in BF16.") - - # max bf16 layers per pipeline stage - max_bf16_layers_per_pipeline_stage = ( - self.num_layers // self.pipeline_model_parallel_size - ) - - # check start/end bf16 layer counts are valid - if self.first_last_layers_bf16: - if ( - self.num_layers_at_start_in_bf16 < 0 - or self.num_layers_at_start_in_bf16 > max_bf16_layers_per_pipeline_stage - ): - raise ValueError( - f"num_layers_at_start_in_bf16 ({self.num_layers_at_start_in_bf16}) must be " - f"between 0 and number of layers per pipeline stage " - f"({max_bf16_layers_per_pipeline_stage})." - ) - if ( - self.num_layers_at_end_in_bf16 < 0 - or self.num_layers_at_end_in_bf16 > max_bf16_layers_per_pipeline_stage - ): - raise ValueError( - f"num_layers_at_end_in_bf16 ({self.num_layers_at_end_in_bf16}) must be " - f"between 0 and number of layers per pipeline stage " - f"({max_bf16_layers_per_pipeline_stage})." - ) - - if self.fp8_recipe == Fp8Recipe.custom: - if not self.fp8_quantizer_factory: - raise ValueError( - "fp8_quantizer_factory must be provided when fp8_recipe is 'custom'. " - "Specify a Python import path (e.g., package.module.quantizer_factory) " - "via --fp8-quantizer-factory." - ) - - if self.fp8_param and not self.fp8: - raise ValueError("fp8_param must be used together with fp8 mode.") - - # FP4 validation - if self.fp4_param and not self.fp4: - raise ValueError("fp4_param must be used together with fp4 mode.") - - if self.fp4 and self.fp8: - raise ValueError("fp4 and fp8 cannot be used simultaneously. Please choose one.") - - if self.fp4 and self.fp4_recipe == Fp4Recipe.custom: - if not self.fp4_quantizer_factory: - raise ValueError( - "fp4_quantizer_factory must be provided when fp4_recipe is 'custom'. " - "Specify a Python import path (e.g., package.module.quantizer_factory) " - "via --fp4-quantizer-factory." - ) - - if self.apply_query_key_layer_scaling: - self.attention_softmax_in_fp32 = True - - if self.expert_model_parallel_size > 1 and self.num_moe_experts is None: - raise ValueError("num_moe_experts must be non None to use expert-parallel.") - - if self.num_moe_experts is not None and self.num_moe_experts <= 0: - raise ValueError("num_moe_experts must be non-negative.") - - if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None: - self.moe_ffn_hidden_size = self.ffn_hidden_size - warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.") - - if self.num_moe_experts is None: - assert ( - self.moe_ffn_hidden_size is None - ), "moe_ffn_hidden_size must be None when num_experts is not set." - - if self.moe_enable_deepep: - if self.moe_token_dispatcher_type != "flex": - raise ValueError("DeepEP backend is only supported with flex token dispatcher.") - if self.moe_flex_dispatcher_backend == "hybridep": - raise ValueError("Only one backend is supported for flex token dispatcher.") - self.moe_flex_dispatcher_backend = "deepep" - warnings.warn( - "moe_enable_deepep is deprecated." - "Please use --moe-flex-dispatcher-backend=deepep instead." - ) - - if self.moe_token_dispatcher_type == "flex": - if self.moe_pad_expert_input_to_capacity and ( - self.moe_enable_deepep or self.moe_flex_dispatcher_backend == "deepep" - ): - raise ValueError( - "Flex token dispatcher with deepep backend does not support " - "moe_pad_expert_input_to_capacity" - ) - - if self.moe_shared_expert_intermediate_size is not None: - if self.moe_shared_expert_intermediate_size <= 0: - raise ValueError( - f"moe_shared_expert_intermediate_size must be " - f"num_shared_experts * ffn_size_of_each_shared_expert, " - f"but got {self.moe_shared_expert_intermediate_size}" - ) - if self.moe_shared_expert_overlap and self.moe_token_dispatcher_type not in [ - "alltoall" - ]: - raise ValueError( - f"moe_shared_expert_overlap only works with alltoall token dispatcher." - ) - - if isinstance(self.moe_router_load_balancing_type, list): - assert isinstance(self.moe_aux_loss_coeff, list) and len( - self.moe_aux_loss_coeff - ) == len(self.moe_router_load_balancing_type), ( - "moe_aux_loss_coeff must be a list of the same length as " - "moe_router_load_balancing_type" - ) - - if self.moe_expert_capacity_factor is not None: - if self.moe_expert_capacity_factor < 0: - self.moe_expert_capacity_factor = None - if isinstance(self.moe_router_load_balancing_type, list): - for load_balancing_type in self.moe_router_load_balancing_type: - if load_balancing_type not in [ - "aux_loss", - "seq_aux_loss", - "global_aux_loss", - "none", - ]: - raise ValueError( - "moe_expert_capacity_factor only works with aux_loss, " - "seq_aux_loss, global_aux_loss or none load balancing" - ) - elif self.moe_router_load_balancing_type not in [ - "aux_loss", - "seq_aux_loss", - "global_aux_loss", - "none", - ]: - raise ValueError( - "moe_expert_capacity_factor only works with aux_loss, " - "seq_aux_loss, global_aux_loss or none load balancing" - ) - - if self.moe_pad_expert_input_to_capacity: - if self.moe_expert_capacity_factor is None: - raise ValueError( - "moe_expert_capacity_factor must be set to use moe_pad_expert_input_to_capacity" - ) - - if self.cpu_offloading and ( - self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers - ): - raise ValueError( - f"CPU offloading can be done only for layers less than {self.num_layers}" - ) - - if self.cpu_offloading and self.pipeline_model_parallel_size > 1: - raise ValueError( - "Currently there is no support for Pipeline parallelism with CPU offloading" - ) - - if self.cpu_offloading and self.recompute_granularity is not None: - raise ValueError( - "CPU offloading does not work when activation recomputation is enabled" - ) - - if self.recompute_granularity is not None: - if self.recompute_granularity not in ["full", "selective"]: - raise ValueError( - f'When using recompute_granuarlity: {self.recompute_granularity} must be "full"' - 'or "selective".' - ) - - if self.recompute_method is not None: - if self.recompute_method not in ["block", "uniform"]: - raise ValueError( - f'recompute_method: {self.recompute_method} must be "block" or "uniform".' - ) - elif self.recompute_granularity != "selective": - raise ValueError( - f"Using recompute_granularity: {self.recompute_granularity} so " - 'recompute_method must be "block" or "uniform"' - ) - - if self.recompute_granularity != "selective" and self.recompute_num_layers is None: - raise ValueError( - f"When using recompute_granularity: {self.recompute_granularity} " - "recompute_num_layers must be between " - "1 and num_layers_per_pipeline_rank: " - f"{self.num_layers // self.pipeline_model_parallel_size}" - ) - elif ( - self.recompute_granularity == "selective" and self.recompute_num_layers is not None - ): - raise ValueError( - f"When using recompute_granularity: {self.recompute_granularity} " - "recompute_num_layers must be None." - ) - - if self.distribute_saved_activations and self.sequence_parallel: - raise ValueError( - f"distribute_saved_activations: {self.distribute_saved_activations} must be " - f"false when sequence parallel is enabled: {self.sequence_parallel}" - ) - - if self.recompute_modules is None: - self.recompute_modules = ["core_attn"] - - if self.recompute_granularity == "selective": - if len(self.recompute_modules) > 0: - allowed_modules = { - "core_attn", - "moe_act", - "layernorm", - "mla_up_proj", - "mlp", - "moe", - "shared_experts", - } - invalid_modules = set(self.recompute_modules) - allowed_modules - assert not invalid_modules, ( - f"Invalid choices for recompute_modules: {invalid_modules}. " - f"Allowed modules are: {allowed_modules}" - ) - - if "moe_act" in self.recompute_modules and not self.moe_grouped_gemm: - raise ValueError( - "moe_act in recompute_modules is only supported with moe_grouped_gemm." - ) - - if "mla_up_proj" in self.recompute_modules and not self.multi_latent_attention: - raise ValueError( - "mla_up_proj in recompute_modules is only supported with " - "multi_latent_attention." - ) - - if "core_attn" in self.recompute_modules: - warnings.warn( - "If you are using transformer_engine as the transformer implementation, " - "the core_attn is from transformer_engine and may be the fused version. " - "For fused attention, you have no need to set 'core_attn' to recompute. " - "Please check that the core_attn recompute is really needed." - ) - - if "shared_experts" in self.recompute_modules: - if ( - self.moe_shared_expert_intermediate_size is not None - and self.moe_shared_expert_overlap - ): - raise ValueError( - "shared_experts recompute cannot work with --moe-shared-expert-overlap." - ) - - if self.fp8: - if "moe_act" in self.recompute_modules or "layernorm" in self.recompute_modules: - if self.fp8_recipe == 'delayed': - raise ValueError( - "Delayed scaling does not support moe_act and layernorm recompute " - "for fp8." - ) - if not is_te_min_version("2.6.0dev0"): - raise ValueError( - "moe_act and layernorm recompute for fp8 needs " - "transformer-engine>=2.6.0dev0, " - f"but your version is {get_te_version()}." - ) - - if self.moe_layer_recompute: - warnings.warn( - "--moe-layer-recompute is deprecated. " - "Use --recompute-granularity selective --recompute-modules moe_layer instead." - ) - if self.recompute_granularity == "full": - raise ValueError( - "Do not set --moe-layer-recompute with full recompute granularity. " - ) - self.recompute_granularity = "selective" - if "moe" not in self.recompute_modules: - self.recompute_modules.append("moe") - - if self.fine_grained_activation_offloading: - assert ( - not self.cpu_offloading - ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading." - assert self.offload_modules is not None and len(self.offload_modules) > 0 - allowed_modules = { - "core_attn", - "attn_proj", - "expert_fc1", - "moe_act", - "attn_norm", - "mlp_norm", - "qkv_linear", - } - invalid_modules = set(self.offload_modules) - allowed_modules - assert not invalid_modules, ( - f'Invalid choices for offload_modules: {invalid_modules}. ' - f'Allowed modules are: {allowed_modules}' - ) - if "attn_proj" in self.offload_modules and "core_attn" not in self.offload_modules: - raise ValueError( - "attn_proj cannot be set to offload_modules alone without core_attn " - "because the input of attn_proj is the output of core_attn, " - "which is needed in core_attn.backward()." - ) - - if ( - self.num_layers_in_first_pipeline_stage is not None - or self.num_layers_in_last_pipeline_stage is not None - ) and ( - self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split - ): - raise ValueError( - "num_layers_in_first_pipeline_stage and num_layers_in_last_pipeline_stage cannot be" - "set at the same time with account_for_embedding_in_pipeline_split" - "and account_for_loss_in_pipeline_split" - ) - - # PP layout - if self.pipeline_model_parallel_layout is not None: - # If pipeline layout is set, we will check the conflicts - # with other pipeline layout arguments. - any_conflict = ( - self.num_layers_in_first_pipeline_stage is not None - or self.num_layers_in_last_pipeline_stage is not None - or self.account_for_embedding_in_pipeline_split - or self.account_for_loss_in_pipeline_split - ) - if any_conflict: - raise ValueError( - "pipeline_model_parallel_layout cannot be set" - " with other pipeline layout arguments." - f" {self.num_layers_in_first_pipeline_stage=}," - f" {self.num_layers_in_last_pipeline_stage=}," - f" {self.account_for_embedding_in_pipeline_split=}," - f" {self.account_for_loss_in_pipeline_split=}." - ) - - # Transfer pipeline_model_parallel_layout from str or list to - # PipelineParallelLayerLayout - if isinstance(self.pipeline_model_parallel_layout, str): - self.pipeline_model_parallel_layout = PipelineParallelLayerLayout.from_str( - layout=self.pipeline_model_parallel_layout, - pipeline_model_parallel_size=self.pipeline_model_parallel_size, - ) - elif isinstance(self.pipeline_model_parallel_layout, list): - # Since list is not hashable, the initialization will not be cached. - self.pipeline_model_parallel_layout = PipelineParallelLayerLayout( - layout=self.pipeline_model_parallel_layout, - pipeline_model_parallel_size=self.pipeline_model_parallel_size, - ) - - # Check whether the input VPP size conflicts with the PP layout - detected_vpp_size = ( - self.pipeline_model_parallel_layout.virtual_pipeline_model_parallel_size - ) - if self.virtual_pipeline_model_parallel_size is not None: - assert self.virtual_pipeline_model_parallel_size == detected_vpp_size, ( - f"virtual_pipeline_model_parallel_size conflicts with" - f" pipeline_model_parallel_layout," - f" ({self.virtual_pipeline_model_parallel_size=}, " - f" {detected_vpp_size=})" - ) - elif detected_vpp_size > 1: - self.virtual_pipeline_model_parallel_size = detected_vpp_size - - # Check whether the layout is valid. - self.mtp_standalone = self.pipeline_model_parallel_layout.validate_layer_layout( - num_layers=self.num_layers, mtp_num_layers=self.mtp_num_layers - ) - - # Uneven PP - elif ( - self.num_layers_in_first_pipeline_stage is not None - or self.num_layers_in_last_pipeline_stage is not None - ): - pipeline_parallel_size = self.pipeline_model_parallel_size - num_layers = self.num_layers - - if self.num_layers_in_first_pipeline_stage is not None: - if self.num_layers_in_first_pipeline_stage <= 0: - raise ValueError("num_layers_in_first_pipeline_stage must be larger than 0") - - if self.virtual_pipeline_model_parallel_size is not None: - if ( - self.num_layers_in_first_pipeline_stage - % self.virtual_pipeline_model_parallel_size - != 0 - ): - raise ValueError( - f"number of layers at first stage: " - f"{self.num_layers_in_first_pipeline_stage}" - f"must be divisible by virtual pipeline" - f"parallel degree {self.virtual_pipeline_model_parallel_size}" - ) - num_layers -= self.num_layers_in_first_pipeline_stage - pipeline_parallel_size -= 1 - - if self.num_layers_in_last_pipeline_stage is not None: - if self.num_layers_in_last_pipeline_stage <= 0: - raise ValueError("num_layers_in_last_pipeline_stage must be larger than 0") - - if self.virtual_pipeline_model_parallel_size is not None: - if ( - self.num_layers_in_last_pipeline_stage - % self.virtual_pipeline_model_parallel_size - != 0 - ): - raise ValueError( - f"number of layers at last stage: " - f"{self.num_layers_in_last_pipeline_stage}" - f"must be divisible by virtual pipeline" - f"parallel degree {self.virtual_pipeline_model_parallel_size}" - ) - num_layers -= self.num_layers_in_last_pipeline_stage - pipeline_parallel_size -= 1 - - # Here pipeline_parallel_size is the number of middle PP stages. If there are middle - # PP stages, check number of layers at middle stage is divisible by middle PP size. - if pipeline_parallel_size and not num_layers % pipeline_parallel_size == 0: - raise ValueError( - f"number of layers at middle stage: {num_layers} must be divisible by" - f"the middle pipeline model parallel size {pipeline_parallel_size}" - ) - - # If there are middle PP stages, check number of layers - # on each middle PP rank is divisible by VPP size. - if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None: - num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size - if ( - not num_layers_per_middle_pipeline_rank - % self.virtual_pipeline_model_parallel_size - == 0 - ): - raise ValueError( - f"number of layers on each middle pipeline rank:" - f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual" - f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}" - ) - - elif ( - self.account_for_embedding_in_pipeline_split or self.account_for_loss_in_pipeline_split - ): - if self.virtual_pipeline_model_parallel_size is None: - num_layers = self.num_layers - - if self.account_for_embedding_in_pipeline_split: - num_layers += 1 - - if self.account_for_loss_in_pipeline_split: - num_layers += 1 - - if not num_layers % self.pipeline_model_parallel_size == 0: - raise ValueError( - f"number of middle layers: {num_layers} must be divisible by " - f"middle pipeline_model_parallel_size {self.pipeline_model_parallel_size}" - ) - else: - num_layers = self.num_layers - if self.account_for_embedding_in_pipeline_split: - num_layers += 1 - - if self.account_for_loss_in_pipeline_split: - num_layers += 1 - - if not num_layers % self.pipeline_model_parallel_size == 0: - raise ValueError( - f"num_layers: {num_layers} after enable" - f"account_for_embedding_in_pipeline_split or " - f"account_for_loss_in_pipeline_split must be divisible" - f"by pipeline_model_parallel_size " - f"{self.pipeline_model_parallel_size}" - ) - - num_layers_per_pipeline_rank = num_layers // self.pipeline_model_parallel_size - if ( - not num_layers_per_pipeline_rank % self.virtual_pipeline_model_parallel_size - == 0 - ): - raise ValueError( - f"number of layers on each pipeline rank: {num_layers_per_pipeline_rank}" - f"(after enable account_for_embedding_in_pipeline_split or " - f"account_for_loss_in_pipeline_split) must be divisible by" - f"virtual_pipeline_model_parallel_size" - f"{self.virtual_pipeline_model_parallel_size}" - ) - - if self.apply_query_key_layer_scaling: - self.attention_softmax_in_fp32 = True - - if self.bias_activation_fusion: - if self.activation_func not in [F.gelu, F.silu, quick_gelu]: - raise ValueError( - "When bias_activation_fusion is True, activation function should be either " - "gelu, swiglu, or quick_geglu" - ) - if ( - self.activation_func == F.gelu - and not self.gated_linear_unit - and not self.add_bias_linear - ): - raise ValueError( - "When bias_activation_fusion is True, gated_linear_unit is False " - "and activation function is gelu, add_bias_linear must also be True." - ) - if self.activation_func == quick_gelu and not self.gated_linear_unit: - raise ValueError( - "When bias_activation_fusion is True and activation function is quick_gelu, " - "gated_linear_unit must be True." - ) - if self.glu_linear_offset != 0.0 and self.activation_func != quick_gelu: - raise ValueError( - "When bias_activation_fusion is True and glu_linear_offset is non-zero, " - "activation function must be quick_gelu." - ) - - if self.use_te_activation_func: - raise ValueError( - "bias_activation_fusion and use_te_activation_func cannot be both true. " - "If you use bias in MLP FC1, we recommend setting bias_activation_fusion " - "to True and use_te_activation_func to False." - ) - - if self.use_te_activation_func: - if self.activation_func not in (F.gelu, F.silu, F.relu): - raise ValueError( - "TransformerEngine only support gelu, geglu, silu, swiglu, relu, reglu. " - "If you don't want to use TransformerEngine activation function, set " - "use_te_activation_func to False" - ) - - if self.activation_func_fp8_input_store: - if self.activation_func != F.silu or not self.gated_linear_unit: - raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.") - - if self.apply_rope_fusion: - if self.multi_latent_attention: - warnings.warn( - "apply_rope_fusion for multi-latent attention only supports training. " - "It is experimental and may change in future versions." - ) - else: - if self.rotary_interleaved: - if not is_te_min_version("2.3.0"): - raise ValueError( - "rotary_interleaved does not work with apply_rope_fusion for " - "TE < 2.3.0. Please install TE >= 2.3.0" - ) - - from megatron.core.models.common.embeddings.rope_utils import ( - fused_apply_rotary_pos_emb, - fused_apply_rotary_pos_emb_thd, - ) - - if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None: - raise ValueError( - "apply_rope_fusion is not available. Please install TE >= 1.4." - ) - - if self.fused_single_qkv_rope: - if self.attention_output_gate: - raise ValueError("fused_single_qkv_rope does not support gated attention for now.") - - if self.multi_latent_attention and self.rotary_interleaved: - raise ValueError("rotary_interleaved does not work with multi_latent_attention.") - - # Set the embedding init method - if self.embedding_init_method_std is None: - # By default, use the same init std as you use for every other non-output layer. - self.embedding_init_method_std = self.init_method_std - - if self.embedding_init_method is None: - if self.init_method is None or (self.embedding_init_method_std != self.init_method_std): - # In this case, we set both the init method and the embedding init method to - # whatever std value requested (or defaulted) for the embedding_init_layer - self.embedding_init_method = init_method_normal(self.embedding_init_method_std) - else: - # Replicate the current behavior where if you are not changing the std of the - # embedding init differently and the init method is set, we fallback to the - # init method for this layer. Since we are here after an OR we know that - # init_method is not None - self.embedding_init_method = self.init_method - - if self.init_method is None: - self.init_method = init_method_normal(self.init_method_std) - - if self.output_layer_init_method is None: - self.output_layer_init_method = scaled_init_method_normal( - self.init_method_std, - self.num_layers, - multiplier=2.0 if not self.is_hybrid_model else 1.0, - ) - - if self.num_moe_experts is not None and self.add_bias_linear: - assert ( - self.expert_tensor_parallel_size == 1 - ), "Bias in Moe is only supported when ETP==1" - - if self.moe_router_enable_expert_bias and self.moe_router_score_function != "sigmoid": - raise ValueError( - "Expert bias for aux-loss-free routing only supports sigmoid score function." - "Please set --moe-router-score-function sigmoid for sigmoid score function." - ) - - if self.num_moe_experts and self.fp8: - # TE version below 1.7.0 will raise Error when handle zeros tokens for expert - if not is_te_min_version("1.7.0.dev0"): - raise ValueError( - "Only transformer-engine>=1.7.0 supports MoE FP8 training, " - f"but your version is {get_te_version()}." - ) - - if self.moe_grouped_gemm and not is_te_min_version("1.11.0"): - raise ValueError( - "Only transformer-engine>=1.11.0 supports FP8 grouped gemm, " - f"but your version is {get_te_version()}." - ) - - if self.moe_router_padding_for_fp8: - # enable moe_router_padding_for_quantization - warnings.warn( - "--moe-router-padding-for-fp8 is going to be deprecated. " - "Use --moe-router-padding-for-quantization instead." - ) - self.moe_router_padding_for_quantization = True - - if self.moe_router_padding_for_quantization: - if self.fp8 is None and self.fp4 is None: - raise ValueError( - "fp8/fp4 must be specified when moe_router_padding_for_quantization is True." - ) - - if self.moe_token_dispatcher_type in ["allgather", "alltoall_seq"]: - raise ValueError( - "allgather and alltoall_seq dispatcher does not support " - "moe_router_padding_for_quantization." - ) - - if ( - self.moe_router_topk == 1 - and self.moe_router_score_function == "softmax" - and not self.moe_router_pre_softmax - and self.moe_router_load_balancing_type != "sinkhorn" - ): - # Requires applying softmax before selecting the top-k when k is 1, - # since softmax on a [num_tokens, 1] would yield a zero gradient. - raise ValueError("Please use --moe-router-pre-softmax when topk is 1.") - - if self.moe_router_group_topk: - if self.moe_router_topk_limited_devices: - raise ValueError( - "moe_router_topk_limited_devices is deprecated and replaced by " - "moe_router_group_topk and moe_router_num_groups." - ) - if not self.moe_router_num_groups: - raise ValueError( - "When using group limited routing, moe_router_num_groups must be specified." - ) - else: - assert self.num_moe_experts % self.moe_router_num_groups == 0, ( - f"num_moe_experts ({self.num_moe_experts}) should be divisible by " - f"moe_router_num_groups ({self.moe_router_num_groups})." - ) - assert self.moe_router_group_topk <= self.moe_router_num_groups, ( - f"moe_router_group_topk ({self.moe_router_group_topk}) should be smaller than " - f"moe_router_num_groups ({self.moe_router_num_groups})." - ) - elif self.moe_router_topk_limited_devices: - warnings.warn( - "moe_router_topk_limited_devices is deprecated. Use moe_router_group_topk and " - "moe_router_num_groups instead." - ) - self.moe_router_group_topk = self.moe_router_topk_limited_devices - self.moe_router_num_groups = self.expert_model_parallel_size - - if self.enable_cuda_graph or self.external_cuda_graph: - assert ( - self.cuda_graph_impl == "none" - ), "Do not use enable_cuda_graph or external_cuda_graph with cuda_graph_impl." - assert ( - not self.enable_cuda_graph or not self.external_cuda_graph - ), "enable_cuda_graph and external_cuda_graph cannot be enabled at the same time." - - if self.enable_cuda_graph: - warnings.warn('enable_cuda_graph is deprecated, use cuda_graph_impl=local instead.') - self.cuda_graph_impl = "local" - if self.external_cuda_graph: - warnings.warn( - 'external_cuda_graph is deprecated, ' - 'use cuda_graph_impl=transformer_engine instead.' - ) - self.cuda_graph_impl = "transformer_engine" - - if self.cuda_graph_scope is None: - self.cuda_graph_scope = [] - elif not isinstance(self.cuda_graph_scope, list): - if isinstance(self.cuda_graph_scope, CudaGraphScope): - self.cuda_graph_scope = [self.cuda_graph_scope] - else: - assert isinstance(self.cuda_graph_scope, str), ( - "cuda_graph_scope must be a string that can be converted to a list of " - f"CudaGraphScope, got {self.cuda_graph_scope}." - ) - self.cuda_graph_scope = self.cuda_graph_scope.split(',') - if all(isinstance(scope, str) for scope in self.cuda_graph_scope): - # Backward compatibility for "full" scope. Now we use an empty list instead. - if "full" in self.cuda_graph_scope: - assert self.cuda_graph_scope == [ - "full" - ], "full scope cannot be used with other scopes." - warnings.warn( - "full scope is deprecated. " - "Use empty cuda_graph_scope to capture the whole layer." - ) - self.cuda_graph_scope = [] - else: - self.cuda_graph_scope = [CudaGraphScope[scope] for scope in self.cuda_graph_scope] - assert all( - isinstance(scope, CudaGraphScope) for scope in self.cuda_graph_scope - ), f"cuda_graph_scope must be a list of CudaGraphScope, got {self.cuda_graph_scope}." - - if self.cuda_graph_impl != "none": - assert self.cuda_graph_impl in [ - "transformer_engine", - "local", - ], f"Invalid cuda graph implementation: {self.cuda_graph_impl}" - - if self.cpu_offloading: - raise ValueError("CUDA graphs not supported with CPU offloading.") - - if self.cuda_graph_impl == "local": - assert not self.cuda_graph_scope or self.cuda_graph_scope == [ - CudaGraphScope.full_iteration - ], ( - "For local cuda graph implementation, the only valid value for " - "cuda_graph_scope is full_iteration, or an empty list to denote layerwise " - "graphs. To use other scopes, use cuda_graph_impl=transformer_engine." - ) - - if self.cuda_graph_impl == "transformer_engine": - assert CudaGraphScope.full_iteration not in self.cuda_graph_scope, ( - "To use full iteration cuda graph, please use " - "cuda_graph_impl=local instead of cuda_graph_impl=transformer_engine." - ) - assert ( - CudaGraphScope.moe not in self.cuda_graph_scope - or CudaGraphScope.moe_router not in self.cuda_graph_scope - ), 'cuda_graph_scope must not contain both moe and moe_router.' - if CudaGraphScope.moe_preprocess in self.cuda_graph_scope: - assert ( - CudaGraphScope.moe_router in self.cuda_graph_scope - ), 'moe_preprocess cuda graph is only supported with moe_router cuda graph.' - if self.num_moe_experts is None or self.num_moe_experts <= 1: - assert ( - CudaGraphScope.moe not in self.cuda_graph_scope - and CudaGraphScope.moe_router not in self.cuda_graph_scope - ), 'moe cuda graph is only supported for MoE.' - else: - if self.moe_layer_freq == 1 or ( - isinstance(self.moe_layer_freq, list) and 0 not in self.moe_layer_freq - ): - assert CudaGraphScope.mlp not in self.cuda_graph_scope, ( - 'mlp cuda graph is only supported for dense layers, ' - 'but not found in the model.' - ) - if ( - self.moe_expert_capacity_factor is None - or not self.moe_pad_expert_input_to_capacity - ): - assert ( - CudaGraphScope.moe not in self.cuda_graph_scope - ), 'moe cuda graph is only supported with drop-padding MoE.' - if self.moe_token_dispatcher_type == 'alltoall' and ( - self.moe_expert_capacity_factor is not None - or self.moe_router_padding_for_quantization - ): - assert CudaGraphScope.moe_preprocess not in self.cuda_graph_scope, ( - 'moe_preprocess cuda graph is not supported when there are ' - 'DtoH copies and synchronizations in the preprocess step.' - ) - - if self.recompute_granularity: - if self.recompute_granularity != "selective" or not self.cuda_graph_scope: - raise ValueError( - "Full-layer CUDA graphs not supported with activation recomputation." - ) - elif self.cuda_graph_scope != [CudaGraphScope.full_iteration]: - # For scoped CUDA graphs, only the non-graphed parts of the layer can be - # recomputed. So check if there are overlaps between the recomputed parts - # and the graphed parts. - if CudaGraphScope.attn in self.cuda_graph_scope: - for module in self.recompute_modules: - if module in ['core_attn', 'mla_up_proj']: - raise ValueError( - f'attn cuda graph is not supported with {module} recompute.' - ) - if ( - CudaGraphScope.mlp in self.cuda_graph_scope - and "mlp" in self.recompute_modules - ): - raise ValueError(f'mlp cuda graph is not supported with mlp recompute.') - if CudaGraphScope.moe in self.cuda_graph_scope: - for module in self.recompute_modules: - if module in ['moe_act', 'moe', 'shared_experts']: - raise ValueError( - f'moe cuda graph is not supported with {module} recompute.' - ) - if CudaGraphScope.moe_router in self.cuda_graph_scope: - for module in self.recompute_modules: - if module in ['moe', 'shared_experts']: - raise ValueError( - f'moe_router cuda graph is not supported with {module} ' - 'recompute.' - ) - if "layernorm" in self.recompute_modules: - if ( - CudaGraphScope.attn in self.cuda_graph_scope - and CudaGraphScope.mlp in self.cuda_graph_scope - and ( - CudaGraphScope.moe in self.cuda_graph_scope - or CudaGraphScope.moe_router in self.cuda_graph_scope - ) - ): - raise ValueError( - 'cuda graph is not supported with layernorm recompute.' - ) - if CudaGraphScope.attn in self.cuda_graph_scope: - warnings.warn( - "input_layernorm recompute is not supported with attention " - "cudagraph. Will only recompute the pre_mlp_layernorm." - ) - if ( - CudaGraphScope.mlp in self.cuda_graph_scope - or CudaGraphScope.moe in self.cuda_graph_scope - or CudaGraphScope.moe_router in self.cuda_graph_scope - ): - warnings.warn( - "pre_mlp_layernorm recompute is not supported with mlp/moe " - "cudagraph. Will only recompute the input_layernorm." - ) - - if self.moe_token_dispatcher_type in ["allgather"]: - if self.variable_seq_lengths is True: - raise ValueError( - f"Token dispatcher type: {self.moe_token_dispatcher_type} does not support " - f"variable sequence length, please use alltoall dispatcher instead." - ) - - if self.moe_permute_fusion: - from megatron.core.transformer.moe.moe_utils import ( - fused_permute, - fused_permute_with_probs, - fused_sort_chunks_by_index, - fused_sort_chunks_by_index_with_probs, - fused_unpermute, - ) - - if ( - fused_permute is None - or fused_permute_with_probs is None - or fused_sort_chunks_by_index is None - or fused_sort_chunks_by_index_with_probs is None - or fused_unpermute is None - ): - raise ValueError("fused permutation is not available. Please install TE >= 2.1.0.") - - if self.overlap_moe_expert_parallel_comm: - # TODO: remove this after we fix the hang issue with torch version < 2.6.0 - assert is_torch_min_version( - "2.6.0" - ), "A2A Overlap encounters hang issue with torch version < 2.6.0" - if self.pipeline_model_parallel_size > 1: - assert self.virtual_pipeline_model_parallel_size is not None, ( - "If enabling EP A2A overlap, virtual_pipeline_model_parallel_size " - "must be specified when pipeline_model_parallel_size > 1" - ) - # Expert model parallelism requirements - assert ( - self.expert_model_parallel_size > 1 - ), 'overlap_moe_expert_parallel_comm is only supported with expert model parallelism' - assert self.moe_token_dispatcher_type in [ - 'alltoall', - 'flex', - ], 'overlap_moe_expert_parallel_comm is supported with alltoall/flex token dispatcher' - - assert ( - self.recompute_granularity != 'full' - ), 'disable full recomputation when enabling overlap_moe_expert_parallel_comm' - assert ( - self.recompute_method is None - ), 'disable recomputation method when enabling overlap_moe_expert_parallel_comm' - assert ( - self.recompute_num_layers is None - ), 'recompute_num_layers must be None when enabling overlap_moe_expert_parallel_comm' - - # Check if bf16 or fp16 is used - assert ( - self.bf16 or self.fp16 - ), 'overlap_moe_expert_parallel_comm is only supported with bf16 or fp16 model' - - assert ( - not self.moe_shared_expert_overlap - ), 'disable moe_shared_expert_overlap when enabling overlap_moe_expert_parallel_comm' - assert ( - self.mtp_num_layers is None or self.mtp_num_layers == 1 - ), 'MTP layernum only supports 1 when enabling overlap_moe_expert_parallel_comm.' - - # Check delay_wgrad_compute compatibility - if self.delay_wgrad_compute: - assert ( - self.overlap_moe_expert_parallel_comm - ), 'overlap_moe_expert_parallel_comm must be enabled when enabling delay_wgrad_compute' - assert ( - not self.moe_use_legacy_grouped_gemm - ), 'delay_wgrad_compute is not supported with legacy groupedgemm implementation' - - if self.context_parallel_size > 1 and self.cp_comm_type is not None: - if isinstance(self.cp_comm_type, list): - assert len(self.cp_comm_type) == self.num_layers, ( - f"Length of cp_comm_type ({len(self.cp_comm_type)}) should equal to " - f"the total number of transformer layers ({self.num_layers})!" - ) - else: - assert isinstance( - self.cp_comm_type, str - ), "Unsupported communication type for context parallelism!" - - assert ( - self.pipeline_model_parallel_size > 0 - ), f"Pipeline model parallel size must be larger than 0 \ - when enable --standalone-embedding-stage and --standalone-loss-stage" - - if ( - self.num_moe_experts is not None - and self.num_moe_experts >= 32 - and not self.moe_router_dtype - ): - warnings.warn( - "Using a large number of experts (e.g. >=32) without fp32 routing. " - "Consider enabling moe_router_dtype for better numerical stability." - ) - if self.symmetric_ar_type is not None: - if not HAVE_PACKAGING: - raise ImportError( - "packaging is not installed. Please install it with `pip install packaging`." - ) - assert is_torch_min_version("2.7.0a0"), "Must have at least torch version 2.7 or higher" - assert is_te_min_version("2.3.0") or get_te_version() == PkgVersion( - "2.3.0.dev0+39c0e70" - ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce" - - if self.no_rope_freq: - assert not self.flash_decode, "flash_decode cannot be used with no_rope." - if isinstance(self.no_rope_freq, int): - assert self.num_layers % self.no_rope_freq == 0, ( - f"no_rope_freq={self.no_rope_freq} should be " - f"divisible by num_layers={self.num_layers}." - ) - # Convert integer pattern to list pattern - # e.g. no_rope=4 with num_layers=8 becomes [0,0,0,1,0,0,0,1] - pattern = [0] * (self.no_rope_freq - 1) + [1] - self.no_rope_freq = pattern * (self.num_layers // self.no_rope_freq) - else: - assert len(self.no_rope_freq) == self.num_layers, ( - f"Length of no_rope list ({len(self.no_rope_freq)}) must match " - f"the number of layers ({self.num_layers})" - ) - - if self.fallback_to_eager_attn: - assert self.transformer_impl == "transformer_engine", ( - f"fallback_to_eager_attn is only available with transformer_engine implementation," - f" but got {self.transformer_impl=}." - ) - - if self.fallback_to_eager_attn or self.transformer_impl == "local": - if self.context_parallel_size > 1 and self.cp_comm_type is not None: - all_cp_comm_types_are_all_gather = ( - all(item == "all_gather" for item in self.cp_comm_type) - if isinstance(self.cp_comm_type, list) - else self.cp_comm_type == "all_gather" - ) - if not all_cp_comm_types_are_all_gather: - raise ValueError( - f"fallback_to_eager_attn only supports all_gather communication type " - f"for context parallelism, but got {self.cp_comm_type=} instead." - ) - - -@dataclass -class MLATransformerConfig(TransformerConfig): - """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers. - - The initialization function has an argument for each parameter, including those in - ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA. - """ - - multi_latent_attention: bool = True - """Whether to use Multi-Latent Attention.""" - - q_lora_rank: int = 512 - """Rank of Query tensor's low rank representation.""" - - kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" - - qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" - - qk_pos_emb_head_dim: int = 64 - """Dimension of the position embedding in the QK projection.""" - - v_head_dim: int = 128 - """Dimension of the head in the V projection.""" - - normalization: str = "RMSNorm" - """Default normalization layer for MLA models is RMSNorm.""" - - rope_type: str = "yarn" - """Type of RoPE to use. Default to yarn, options are rope and yarn.""" - - rotary_base: float = 10000 - """Rotary base for the rotary embeddings, used by rope and yarn.""" - - rotary_percent: float = 1.0 - """Rotary percent for the rotary embeddings, used by rope.""" - - rotary_scaling_factor: float = 40 - """Rotary scaling factor for the rotary embeddings, used by yarn.""" - - original_max_position_embeddings: int = 4096 - """Original maximum position embeddings for the original model, used by yarn.""" - - beta_fast: float = 32 - """Beta fast for YaRN RoPE, used by yarn.""" - - beta_slow: float = 1 - """Beta slow for YaRN RoPE, used by yarn.""" - - mscale: float = 1.0 - """Mscale for YaRN RoPE in Multi-Latent Attention, used by yarn.""" - - mscale_all_dim: float = 0.0 - """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" - - cache_mla_latents: bool = False - """Cache the low dimensional tensors for MLA rather than full KV cache. - This is only for the dynamic inference backend and requires that - Flash MLA is installed.""" - - def __post_init__(self): - super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": - raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") - - if self.attention_output_gate: - raise NotImplementedError("Output gate is not supported for MLA yet.") - - if self.cache_mla_latents: - assert ( - self.apply_rope_fusion is False - ), "Rope Fusion is not compatible with caching latents" \ No newline at end of file diff --git a/docker/deepseekv32/megatron.patch b/docker/deepseekv32/megatron.patch index bd5ef87c66c..f7204197c3a 100644 --- a/docker/deepseekv32/megatron.patch +++ b/docker/deepseekv32/megatron.patch @@ -798,6 +798,34 @@ index 353b31e9b..221e93500 100644 + output = unfused_dsa_fn_with_cp(query, key, dim_v, topk_indices, self.softmax_scale) + return output +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index 28cff06f5..befb5c124 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -586,6 +586,9 @@ def topk_routing_with_score_function( + ) + else: + return torch.topk(scores, k=topk, dim=1) ++ ++ from miles.utils.routing_replay import get_routing_replay_compute_topk ++ compute_topk = get_routing_replay_compute_topk(compute_topk) + + if score_function == "softmax": + if use_pre_softmax: +diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py +index 16fc9d9af..3c50a9516 100644 +--- a/megatron/core/transformer/moe/router.py ++++ b/megatron/core/transformer/moe/router.py +@@ -200,6 +200,9 @@ class TopKRouter(Router): + else: + self.global_tokens_per_expert = None + self.ga_steps = None ++ ++ from miles.utils.routing_replay import register_routing_replay ++ register_routing_replay(self) + + def _maintain_float32_expert_bias(self): + """ diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index ed90fdffa..7a7597d66 100644 --- a/megatron/core/transformer/multi_latent_attention.py diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 95803c73f15..4bb7548ea75 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -464,7 +464,7 @@ def update_weights(self) -> None: if isinstance(num_new_engines, tuple): num_new_engines = num_new_engines[0] - + if num_new_engines > 0: self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock) dist.barrier(group=get_gloo_group()) diff --git a/miles/utils/data.py b/miles/utils/data.py index eb512e51481..d258b4070c1 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -209,8 +209,9 @@ def __init__( add_generation_prompt=True, **apply_chat_template_kwargs, ) - except Exception as e: + except Exception: from sglang.srt.entrypoints.openai.encoding_dsv32 import encode_messages + encode_config = dict(thinking_mode="thinking", drop_thinking=True, add_default_bos_token=True) prompt = encode_messages(prompt, **encode_config) ### DSV32 diff --git a/scripts/run_deepseek_v32.py b/scripts/run_deepseek_v32.py index 6d76ade8611..3d46d0e5172 100644 --- a/scripts/run_deepseek_v32.py +++ b/scripts/run_deepseek_v32.py @@ -25,9 +25,9 @@ class ScriptArgs(U.ExecuteTrainConfig): task: Literal["dapo_aime", "gsm8k"] = "dapo_aime" enable_deepep: bool = True data_dir: str = "/root" - model_dir: str = "/root/.cache/dsv32" - model_local_dir: str = "/root/.cache/dsv32" - save_dir: str = "/root/.cache/dsv32" + model_dir: str = "/root/models" + model_local_dir: str = "/root/models" + save_dir: str = "/root/models" megatron_path: str = "/root/Megatron-LM" @@ -115,11 +115,11 @@ def train(args: ScriptArgs): "--rollout-shuffle " "--rm-type math " "--num-rollout 3000 " - "--rollout-batch-size 8 " - "--n-samples-per-prompt 8 " + "--rollout-batch-size 1 " + "--n-samples-per-prompt 1 " "--rollout-temperature 0.8 " # ------------ - "--num-steps-per-rollout 4 " + "--num-steps-per-rollout 1 " "--balance-data " ) @@ -212,6 +212,8 @@ def train(args: ScriptArgs): "--entropy-coef 0.00 " "--eps-clip 0.2 " "--eps-clip-high 0.28 " + "--use-miles-router " + "--use-rollout-routing-replay " ) optimizer_args = ( @@ -246,6 +248,7 @@ def train(args: ScriptArgs): f"--sglang-max-running-requests {sglang_world_size * sglang_decode_max_bs // sglang_attn_tp_size} " f"--sglang-chunked-prefill-size {sglang_world_size * sglang_decode_max_bs} " f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} " + "--sglang-disable-cuda-graph " # For quick experiments # """--sglang-json-model-override-args '{"num_hidden_layers": 5}' """ )