diff --git a/.gitleaks.toml b/.gitleaks.toml index f49343e18..8f2d866a7 100644 --- a/.gitleaks.toml +++ b/.gitleaks.toml @@ -70,13 +70,15 @@ tags = ["internal", "token"] [[rules]] id = "private-ip-10-range" description = "Detects hardcoded private IP addresses in 10.0.0.0/8 range (RFC 1918)" -regex = '''\b10\.\d{1,3}\.\d{1,3}\.\d{1,3}\b''' +regex = '''\b10\.\d{1,3}\.\d{1,3}\.\d{1,3}\b(?:/\d{1,2}\b)?''' tags = ["internal", "ip", "private"] [rules.allowlist] regexTarget = "match" regexes = [ # 10.0.0.1 is a conventional placeholder IP in examples and tests. '''^10\.0\.0\.1$''', + # The full RFC 1918 network is public, not a deployment-specific endpoint. + '''^10\.0\.0\.0/8$''', ] paths = ['''scripts/ci/benchmark\.sh$'''] diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 38cc13649..8e0568fe6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -50,7 +50,7 @@ repos: hooks: - id: docformatter name: docformatter - entry: docformatter --in-place --wrap-descriptions 79 + entry: python .pre-commit-hooks/docformatter_compat.py --in-place --wrap-descriptions 79 language: python types: [python] additional_dependencies: ["docformatter==1.3.1"] diff --git a/.pre-commit-hooks/docformatter_compat.py b/.pre-commit-hooks/docformatter_compat.py new file mode 100644 index 000000000..0d142a0cd --- /dev/null +++ b/.pre-commit-hooks/docformatter_compat.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Run docformatter with correct multiline string token positions.""" + +from __future__ import annotations + +import tokenize +from collections.abc import Callable, Iterator + + +_original_generate_tokens = tokenize.generate_tokens + + +def _generate_tokens(readline: Callable[[], str]) -> Iterator[tokenize.TokenInfo]: + for token in _original_generate_tokens(readline): + if token.type == tokenize.STRING and token.start[0] != token.end[0]: + # Python 3.12.0 can undercount this column after non-ASCII text. + # untokenize otherwise copies the apparent gap after the string, + # duplicating source characters, including its closing quotes. + end_column = len(token.string.rsplit("\n", 1)[-1]) + token = token._replace(end=(token.end[0], end_column)) + yield token + + +def main() -> int: + import docformatter + + original = tokenize.generate_tokens + tokenize.generate_tokens = _generate_tokens + try: + return docformatter.main() + finally: + tokenize.generate_tokens = original + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.pre-commit-hooks/gitleaks_tracked.py b/.pre-commit-hooks/gitleaks_tracked.py index 82d40985d..97126c2a0 100644 --- a/.pre-commit-hooks/gitleaks_tracked.py +++ b/.pre-commit-hooks/gitleaks_tracked.py @@ -90,6 +90,8 @@ def main(argv: Sequence[str] | None = None) -> int: stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + encoding="utf-8", + errors="replace", ) _relay(result.stdout, sys.stdout, snapshot_root) diff --git a/AGENTS.md b/AGENTS.md index beb05ef56..cfbb8e3f1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,20 +8,36 @@ Relax 是一个基于 Ray Serve 的大模型强化学习训练框架，支持 Me ``` relax/ 核心框架 +├── agentic/ Agentic 层 — Session、pipeline、agent process 与训练导出 ├── core/ 编排层 — 训练循环、服务基类、全局注册表 ├── components/ 组件层 — RL 服务组件（Ray Serve Deployment） ├── engine/ 引擎层 — Rollout 数据生成、奖励计算、请求路由 ├── backends/ 后端层 — Megatron 训练后端、SGLang 推理引擎 ├── distributed/ 分布式层 — Ray 集群管理、分布式 Checkpoint ├── entrypoints/ 入口层 — 训练入口脚本 +├── models/ 模型层 — 模型专属实现与注册 └── utils/ 基础设施 — 工具函数、指标监控、多模态处理 -tests/ 测试（镜像 relax/ 层级） +tests/ 模块测试与集成回归测试 ├── backends/megatron/ Megatron 后端测试（权重转换等） +├── backends/sglang/ SGLang 后端测试 +├── components/ RL 服务组件测试 +├── core/ Controller、Service 与注册表测试 +├── data/ 数据处理与 SFT 数据测试 +├── distributed/checkpoint_service/ 分布式 Checkpoint 测试 ├── distributed/ray/ 分布式 / Ray 测试（弹性伸缩等） ├── engine/rewards/ 奖励函数测试 ├── engine/rollout/ Rollout 引擎测试（预取、数据源等） -└── utils/ 工具函数测试（HTTP、指标、流式数据集等） -transfer_queue/ 分布式数据传输队列 +├── engine/sft/ SFT 引擎测试 +├── entrypoints/ 入口行为测试 +├── examples/ 示例级回归测试 +├── integration/ 跨模块集成测试 +├── models/ 模型专属测试 +├── tools/ 工具与辅助脚本测试 +├── utils/ 工具函数测试（HTTP、指标、流式数据集等） +└── test_agentic_rollout.py Agentic runtime 与 Session 测试 +docs/ 中英文用户文档 +skills/ 仓库开发与运维工作流 +docker/ 训练镜像与依赖 patch examples/ 用户级示例（deepeyes、OPD 等） scripts/ 训练启动脚本 & 模型配置 configs/env.yaml 运行时环境配置 diff --git a/docker/Dockerfile b/docker/Dockerfile index 70e98cd9f..1c4a7b792 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,7 +1,7 @@ ARG HTTP_PROXY ARG HTTPS_PROXY ARG NO_PROXY -ARG BASE_IMAGE=mirror.ccs.tencentyun.com/lmsysorg/sglang:v0.5.15.post1-cu129 +ARG BASE_IMAGE=mirror.ccs.tencentyun.com/lmsysorg/sglang:v0.5.17-cu129 ARG TRAIN_IMAGE=train FROM ${BASE_IMAGE} as base @@ -38,9 +38,24 @@ RUN pip install nvidia-cudnn-cu12==9.16.0.29 FROM base as train RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir && \ - pip install --no-cache-dir flash-linear-attention==0.4.1 && \ + pip install --no-cache-dir flash-linear-attention==0.4.2 && \ pip install --no-cache-dir tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ +# DeepSeek-V4 DSA deps, both listed under mcore's `no_pypi_wheels` so built from pinned source: +# - fast-hadamard-transform: required by the Lightning Indexer (dsa.rotate_activation asserts on it) +# - FlashMLA: fused DSA sparse attention; without it the CSA backward is quadratic in seq len +# (mcore only allows the fusion on SM90 when dsa_indexer_loss_coeff == 0) +ARG FLASH_MLA_COMMIT=b7643bd54521f563b839b98289b5cd048c062ba2 +RUN MAX_JOBS=64 pip install --no-cache-dir --no-build-isolation --no-deps \ + git+https://github.com/Dao-AILab/fast-hadamard-transform.git@f134af63deb2df17e1171a9ec1ea4a7d8604d5ca && \ + git clone --recurse-submodules https://github.com/deepseek-ai/FlashMLA.git /opt/FlashMLA && \ + cd /opt/FlashMLA && git checkout ${FLASH_MLA_COMMIT} && \ + git submodule update --init --recursive && \ + MAX_JOBS=64 pip install --no-cache-dir --no-build-isolation --no-deps . && \ + cd / && python -c "import flash_mla; assert hasattr(flash_mla, 'flash_mla_sparse_fwd'), \ + 'flash_mla built but flash_mla_sparse_fwd missing'" && \ + rm -rf /opt/FlashMLA + # FA3 (Hopper flash-attention), built from source. This commit's _flash_attn_forward carries # window_size_left/right to match TE 2.14.1 (docs/draft/sglang-0.5.12-upgrade-plan.md §8.4). # The `cp` exposes flash_attn_3.flash_attn_interface, which is what TE imports. BUT this commit's @@ -70,6 +85,7 @@ RUN MAX_JOBS=64 \ RUN pip -v install --no-cache-dir --no-build-isolation "transformer_engine[pytorch]==2.14.1" && \ TMS_CUDA_MAJOR=$(python -c 'import torch; print(torch.version.cuda.split(".")[0])') pip install git+https://github.com/redai-studio/torch_memory_saver.git@afc13785c50119048e2dd8ac497cc9e29ec75bd4 --no-cache-dir --force-reinstall && \ pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation --no-cache-dir && \ + pip install megatron-energon fsspec==2024.3.1 --no-cache-dir&& \ pip install "numpy<2" nvidia-cudnn-cu12==9.16.0.29 --no-cache-dir && \ NVCC_APPEND_FLAGS="--threads 32" \ pip -v install --disable-pip-version-check --no-cache-dir \ @@ -93,7 +109,7 @@ WORKDIR /root ARG PATCH_VERSION=latest ARG ENABLE_SGLANG_PATCH=1 -ARG MEGATRON_BRIDGE_COMMIT=2faedbf6fe3c422835a44b2b360cadcb2a116a54 +ARG MEGATRON_BRIDGE_COMMIT=af17edf52c514c58c79cd291b04a9b42bc5c57f3 ENV MEGATRON_BRIDGE_COMMIT=${MEGATRON_BRIDGE_COMMIT} \ PYTHONPATH=/root/Megatron-LM/ diff --git a/docker/patch/latest/megatron.patch b/docker/patch/latest/megatron.patch index ec9557dc5..c7ffa0d11 120000 --- a/docker/patch/latest/megatron.patch +++ b/docker/patch/latest/megatron.patch @@ -1 +1 @@ -../megatron/20260506-85bced0ae.patch \ No newline at end of file +../megatron/20260728-0e6ac576f.patch \ No newline at end of file diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index 33c1dc275..5cc532326 120000 --- a/docker/patch/latest/sglang.patch +++ b/docker/patch/latest/sglang.patch @@ -1 +1 @@ -../sglang/v0.5.15.post1.patch \ No newline at end of file +../sglang/v0.5.17.patch \ No newline at end of file diff --git a/docker/patch/megatron/20260728-0e6ac576f.patch b/docker/patch/megatron/20260728-0e6ac576f.patch new file mode 100644 index 000000000..9a3404d19 --- /dev/null +++ b/docker/patch/megatron/20260728-0e6ac576f.patch @@ -0,0 +1,1866 @@ +diff --git a/megatron/bridge/models/gemma/gemma4_bridge.py b/megatron/bridge/models/gemma/gemma4_bridge.py +index a62eb81..c858022 100644 +--- a/megatron/bridge/models/gemma/gemma4_bridge.py ++++ b/megatron/bridge/models/gemma/gemma4_bridge.py +@@ -499,71 +499,81 @@ class Gemma4Bridge(MegatronModelBridge): + """Text-only CausalLM: weights at ``model.*``; override in VL subclass.""" + return "model." + +- def _moe_mapping_registry(self) -> MegatronMappingRegistry: +- """Parameter mappings for the MoE variant.""" ++ def _moe_mapping_registry(self, megatron_prefix: str = "") -> MegatronMappingRegistry: ++ """Parameter mappings for the MoE variant. ++ ++ ``megatron_prefix`` and :meth:`_hf_layer_prefix` mirror ++ :meth:`_dense_mapping_registry`, so a VL subclass can reuse this for a ++ text-only conversion of a ``Gemma4ForConditionalGeneration`` checkpoint, ++ whose language weights sit under ``model.language_model.*``. ++ """ ++ mp = megatron_prefix ++ hp = self._hf_layer_prefix() + param_mappings = { +- "embedding.word_embeddings.weight": "model.embed_tokens.weight", +- "decoder.final_layernorm.weight": "model.norm.weight", +- "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", +- "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", +- "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", +- "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", +- "decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": ( +- "model.layers.*.post_attention_layernorm.weight" ++ f"{mp}embedding.word_embeddings.weight": f"{hp}embed_tokens.weight", ++ f"{mp}decoder.final_layernorm.weight": f"{hp}norm.weight", ++ f"{mp}decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": ( ++ f"{hp}layers.*.input_layernorm.weight" ++ ), ++ f"{mp}decoder.layers.*.self_attention.q_layernorm.weight": f"{hp}layers.*.self_attn.q_norm.weight", ++ f"{mp}decoder.layers.*.self_attention.k_layernorm.weight": f"{hp}layers.*.self_attn.k_norm.weight", ++ f"{mp}decoder.layers.*.self_attention.linear_proj.weight": f"{hp}layers.*.self_attn.o_proj.weight", ++ f"{mp}decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": ( ++ f"{hp}layers.*.post_attention_layernorm.weight" + ), +- "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.pre_feedforward_layernorm_2.weight", +- "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", +- "decoder.layers.*.mlp.post_shared_expert_layernorm.weight": ( +- "model.layers.*.post_feedforward_layernorm_1.weight" ++ f"{mp}decoder.layers.*.pre_mlp_layernorm.weight": f"{hp}layers.*.pre_feedforward_layernorm_2.weight", ++ f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc2.weight": f"{hp}layers.*.mlp.down_proj.weight", ++ f"{mp}decoder.layers.*.mlp.post_shared_expert_layernorm.weight": ( ++ f"{hp}layers.*.post_feedforward_layernorm_1.weight" + ), +- "decoder.layers.*.mlp.router.weight": "model.layers.*.router.proj.weight", ++ f"{mp}decoder.layers.*.mlp.router.weight": f"{hp}layers.*.router.proj.weight", + } + + mapping_list = [AutoMapping(megatron_param=m, hf_param=h) for m, h in param_mappings.items()] + mapping_list.extend( + [ + _Gemma4QKVMapping( +- megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", +- q="model.layers.*.self_attn.q_proj.weight", +- k="model.layers.*.self_attn.k_proj.weight", +- v="model.layers.*.self_attn.v_proj.weight", ++ megatron_param=f"{mp}decoder.layers.*.self_attention.linear_qkv.weight", ++ q=f"{hp}layers.*.self_attn.q_proj.weight", ++ k=f"{hp}layers.*.self_attn.k_proj.weight", ++ v=f"{hp}layers.*.self_attn.v_proj.weight", + ), + GatedMLPMapping( +- megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", +- gate="model.layers.*.mlp.gate_proj.weight", +- up="model.layers.*.mlp.up_proj.weight", ++ megatron_param=f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc1.weight", ++ gate=f"{hp}layers.*.mlp.gate_proj.weight", ++ up=f"{hp}layers.*.mlp.up_proj.weight", + ), + FusedGatedExpertMapping( +- megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", +- hf_param="model.layers.*.experts.gate_up_proj", ++ megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc1.weight*", ++ hf_param=f"{hp}layers.*.experts.gate_up_proj", + ), + FusedExpertMapping( +- megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", +- hf_param="model.layers.*.experts.down_proj", ++ megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc2.weight*", ++ hf_param=f"{hp}layers.*.experts.down_proj", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.layer_scalar", +- hf_param="model.layers.*.layer_scalar", ++ megatron_param=f"{mp}decoder.layers.*.layer_scalar", ++ hf_param=f"{hp}layers.*.layer_scalar", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.mlp.router.per_expert_scale", +- hf_param="model.layers.*.router.per_expert_scale", ++ megatron_param=f"{mp}decoder.layers.*.mlp.router.per_expert_scale", ++ hf_param=f"{hp}layers.*.router.per_expert_scale", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.mlp.router.scale", +- hf_param="model.layers.*.router.scale", ++ megatron_param=f"{mp}decoder.layers.*.mlp.router.scale", ++ hf_param=f"{hp}layers.*.router.scale", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.pre_shared_expert_layernorm.weight", +- hf_param="model.layers.*.pre_feedforward_layernorm.weight", ++ megatron_param=f"{mp}decoder.layers.*.pre_shared_expert_layernorm.weight", ++ hf_param=f"{hp}layers.*.pre_feedforward_layernorm.weight", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.mlp.post_moe_layernorm.weight", +- hf_param="model.layers.*.post_feedforward_layernorm_2.weight", ++ megatron_param=f"{mp}decoder.layers.*.mlp.post_moe_layernorm.weight", ++ hf_param=f"{hp}layers.*.post_feedforward_layernorm_2.weight", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.post_ffn_layernorm.weight", +- hf_param="model.layers.*.post_feedforward_layernorm.weight", ++ megatron_param=f"{mp}decoder.layers.*.post_ffn_layernorm.weight", ++ hf_param=f"{hp}layers.*.post_feedforward_layernorm.weight", + ), + ] + ) +diff --git a/megatron/bridge/models/gemma/modeling_gemma4.py b/megatron/bridge/models/gemma/modeling_gemma4.py +index 1c1dda9..b5a0a9b 100644 +--- a/megatron/bridge/models/gemma/modeling_gemma4.py ++++ b/megatron/bridge/models/gemma/modeling_gemma4.py +@@ -456,13 +456,27 @@ class Gemma4DenseSelfAttention(SelfAttention): + self, + hidden_states: Tensor, + key_value_states=None, +- output_gate: bool = False, +- split_qkv: bool = True, ++ **kwargs, + ): ++ """Take core's qkv keyword arguments verbatim and forward them unchanged. ++ ++ Enumerating them instead pins this class to one Megatron-LM lineage: the ++ .dev.commit core (0e6ac576) added `head_wise_gate` to ++ Attention.get_query_key_value_tensors, between `output_gate` and ++ `split_qkv`, and Attention.forward passes it unconditionally by keyword. ++ Gemma4SelfAttention below already takes **kwargs for the same reason. ++ """ ++ split_qkv = kwargs.get("split_qkv", True) ++ output_gate = kwargs.get("output_gate", False) ++ + if self.is_kv_shared_layer: + if not split_qkv or output_gate: +- return super().get_query_key_value_tensors(hidden_states, key_value_states, output_gate, split_qkv) +- query, _k, _v = super().get_query_key_value_tensors(hidden_states, key_value_states, False, True) ++ return super().get_query_key_value_tensors(hidden_states, key_value_states, **kwargs) ++ # Not forwarding **kwargs here is deliberate: this path unpacks a ++ # plain (q, k, v), and any gate flag would make core return more. ++ query, _k, _v = super().get_query_key_value_tensors( ++ hidden_states, key_value_states, output_gate=False, split_qkv=True ++ ) + kv_source = self._kv_source_ref() if self._kv_source_ref is not None else None + if kv_source is not None and kv_source._stored_kv is not None: + key, value = kv_source._stored_kv +@@ -479,7 +493,7 @@ class Gemma4DenseSelfAttention(SelfAttention): + key_value_states, + ) + else: +- result = super().get_query_key_value_tensors(hidden_states, key_value_states, output_gate, split_qkv) ++ result = super().get_query_key_value_tensors(hidden_states, key_value_states, **kwargs) + if not split_qkv: + return result + if output_gate: +diff --git a/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py b/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py +index 918cac7..2bd3011 100644 +--- a/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py ++++ b/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py +@@ -92,6 +92,15 @@ class Gemma4VLBridge(Gemma4Bridge): + + self._is_dense = False + ++ if self._conversion_mode() == "text": ++ # Same shape as the dense branch above: drop the vision/audio towers and ++ # build the plain MoE GPT provider from ``text_config``. Every field ++ # ``_build_moe_provider`` reads (sliding_window, rope_parameters, head_dim, ++ # global_head_dim, num_global_key_value_heads, attention_k_eq_v, ++ # layer_types, num_experts, top_k_experts, moe_intermediate_size, ++ # intermediate_size, final_logit_softcapping) lives on ``text_config``. ++ return self._build_moe_provider(text_config) ++ + provider_kwargs = self.hf_config_to_provider_kwargs(text_config) + provider = Gemma4VLModelProvider(**provider_kwargs) + +@@ -226,6 +235,10 @@ class Gemma4VLBridge(Gemma4Bridge): + if self._conversion_mode() == "text": + return self._dense_mapping_registry(megatron_prefix="") + return self._dense_vl_mapping_registry() ++ if self._conversion_mode() == "text": ++ # ``_hf_layer_prefix`` still resolves to ``model.language_model.`` here -- ++ # the checkpoint layout does not change, only which submodules we build. ++ return self._moe_mapping_registry(megatron_prefix="") + return self._moe_vl_mapping_registry() + + def _dense_vl_mapping_registry(self) -> MegatronMappingRegistry: +diff --git a/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py b/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +index fa13752..b6c0cb9 100644 +--- a/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py ++++ b/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +@@ -191,8 +191,15 @@ class KimiK25VLBridge(MegatronModelBridge): + for fqn, tensor in converted_weights_dict.items(): + if self._is_quantized_expert_key(fqn): + base = fqn[:-7] if fqn.endswith(".weight") else fqn +- # Preserve the original scale dtype from the HF checkpoint + orig_scale_key = f"{base}.weight_scale" ++ # When the source HF checkpoint has been pre-cast to BF16 (no ++ # `weight_scale` triplet present), passthrough instead of re- ++ # quantizing — downstream sglang loads BF16 and would reject ++ # the INT4 export names with "not found in params_dict". ++ if orig_scale_key not in hf_state_dict: ++ result[fqn] = tensor ++ continue ++ # Preserve the original scale dtype from the HF checkpoint + scale_dtype = ( + hf_state_dict[orig_scale_key].dtype if orig_scale_key in hf_state_dict else torch.bfloat16 + ) +diff --git a/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py b/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py +index ba632f3..99e8aa9 100644 +--- a/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py ++++ b/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py +@@ -60,6 +60,11 @@ class KimiK25VLModelProvider(MLAModelProvider): + pad_token_id: int = 163839 + ignore_index: int = -100 + ++ # Split vision encoder workload across TP ranks (data-parallel over TP). ++ # Each TP rank processes a chunk of images, then all-reduce gathers the ++ # full embedding. Reduces per-GPU peak memory for the vision encoder. ++ vision_dp_when_tp: bool = False ++ + # Freeze options for fine-tuning scenarios + freeze_language_model: bool = False + freeze_vision_model: bool = False +diff --git a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py +index c2bf410..073f570 100644 +--- a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py ++++ b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py +@@ -16,7 +16,9 @@ import logging + from typing import List, Optional + + import torch ++import torch.distributed + from megatron.core import parallel_state ++from megatron.core import parallel_state as mpu + from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region + from megatron.core.transformer.module import MegatronModule +@@ -25,6 +27,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module + from transformers.utils import is_flash_attn_2_available + + from megatron.bridge.models.gpt_provider import GPTModelProvider ++from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.utils import preprocess_packed_seqs + from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync + + +@@ -388,8 +391,79 @@ class KimiK25VLModel(MegatronModule): + + return final_embedding, final_attention_mask, final_labels, position_ids + ++ def _vision_forward_tp_split( ++ self, ++ pixel_values: torch.Tensor, ++ grid_thws: torch.Tensor, ++ ) -> List[torch.Tensor]: ++ """Run vision encoder + projector with workload split across TP ranks. ++ ++ Each TP rank processes a subset of images determined by splitting ++ ``grid_thws``, then the partial feature tensors are all-reduced so ++ every rank holds the complete result. ++ """ ++ tp_rank = mpu.get_tensor_model_parallel_rank() ++ tp_size = mpu.get_tensor_model_parallel_world_size() ++ ++ num_images = grid_thws.shape[0] ++ merge_h, merge_w = self.vision_tower.merge_kernel_size ++ param_dtype = next(self.vision_tower.parameters()).dtype ++ text_hidden = self.projector_config.hidden_size ++ ++ pixel_counts = grid_thws.prod(dim=-1) ++ out_token_counts = (grid_thws[:, 1] // merge_h) * (grid_thws[:, 2] // merge_w) ++ total_out_tokens = out_token_counts.sum().item() ++ ++ chunk_indices = list(range(num_images)) ++ chunks = [chunk_indices[i::tp_size] for i in range(tp_size)] ++ my_indices = chunks[tp_rank] if tp_rank < len(chunks) else [] ++ ++ out_buffer = torch.zeros( ++ (total_out_tokens, text_hidden), ++ device=pixel_values.device, ++ dtype=param_dtype, ++ ) ++ ++ if my_indices: ++ pixel_cumsum = pixel_counts.cumsum(dim=0) ++ pv_parts = [] ++ grid_parts = [] ++ for idx in my_indices: ++ px_start = 0 if idx == 0 else pixel_cumsum[idx - 1].item() ++ px_end = pixel_cumsum[idx].item() ++ pv_parts.append(pixel_values[px_start:px_end]) ++ grid_parts.append(grid_thws[idx : idx + 1]) ++ ++ local_pv = torch.cat(pv_parts, dim=0) ++ local_grid = torch.cat(grid_parts, dim=0) ++ ++ local_vit_out = self.vision_tower(local_pv, local_grid) ++ local_features = self.mm_projector(local_vit_out) ++ ++ out_cumsum = out_token_counts.cumsum(dim=0) ++ for feat_i, img_idx in enumerate(my_indices): ++ out_start = 0 if img_idx == 0 else out_cumsum[img_idx - 1].item() ++ out_end = out_cumsum[img_idx].item() ++ out_buffer[out_start:out_end] = local_features[feat_i].to(param_dtype) ++ ++ tp_group = mpu.get_tensor_model_parallel_group() ++ torch.distributed.all_reduce(out_buffer, group=tp_group) ++ ++ result = [] ++ out_cumsum = out_token_counts.cumsum(dim=0) ++ for i in range(num_images): ++ start = 0 if i == 0 else out_cumsum[i - 1].item() ++ end = out_cumsum[i].item() ++ result.append(out_buffer[start:end]) ++ ++ return result ++ + def _extract_image_features(self, pixel_values, grid_thws): + """Extract and project image features.""" ++ tp_size = mpu.get_tensor_model_parallel_world_size() ++ if getattr(self.config, "vision_dp_when_tp", False) and tp_size > 1: ++ return self._vision_forward_tp_split(pixel_values, grid_thws) ++ + image_features = self.vision_tower(pixel_values, grid_thws) + return self.mm_projector(image_features) + +@@ -430,6 +504,12 @@ class KimiK25VLModel(MegatronModule): + cp_size = parallel_state.get_context_parallel_world_size() + cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 + if self.pre_process: ++ # Save the caller-supplied per-sample attention mask before any rewrite — ++ # _merge_input_ids_with_image_features sets `attention_mask = None` on the ++ # vision path, but the THD repack below needs the original [B, T] mask to ++ # know each sample's valid length. ++ saved_attention_mask = attention_mask ++ + if inputs_embeds is None: + inputs_embeds = self.language_model.embedding( + input_ids=input_ids, position_ids=None +@@ -465,8 +545,30 @@ class KimiK25VLModel(MegatronModule): + # Don't need attention mask for causal attention. + attention_mask = None + +- # Transpose back to (T, B, D) for Megatron language model +- inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) ++ # When THD packed_seq_params is provided (VL+CP/SP), repack the raw ++ # padded [B, T_max, D] embedding into compact THD ++ # [sum(padded_seqlens)/cp, 1, D] using the saved per-sample attention ++ # mask. Mirrors the Qwen3VL bridge: without this, downstream MLA ++ # attention sees a tensor whose first dim does not match ++ # cu_seqlens_q_padded (sum=sum(padded_seqlens)), and SP scatter can ++ # hit `T_max % tp_size != 0` since T_max is raw batch-max. ++ # preprocess_packed_seqs also recomputes cu_seqlens with align64. ++ needs_thd_repack = ( ++ packed_seq_params is not None ++ and packed_seq_params.qkv_format == "thd" ++ and saved_attention_mask is not None ++ ) ++ if needs_thd_repack: ++ inputs_embeds, packed_seq_params = preprocess_packed_seqs( ++ inputs_embeds, # [B, T_max, D] ++ saved_attention_mask, ++ pre_process=True, ++ ) ++ # preprocess_packed_seqs returns [1, T_thd, D]; switch to (T_thd, 1, D) ++ inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() ++ else: ++ # Transpose back to (T, B, D) for Megatron language model ++ inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) + + if cp_size > 1: + inputs_embeds = _split_on_cp_rank(inputs_embeds, cp_size, cp_rank, seq_dim=0) +diff --git a/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/megatron/bridge/models/qwen/qwen3_moe_bridge.py +index a6476f4..86d83b3 100755 +--- a/megatron/bridge/models/qwen/qwen3_moe_bridge.py ++++ b/megatron/bridge/models/qwen/qwen3_moe_bridge.py +@@ -69,6 +69,40 @@ class Qwen3MoEBridge(MegatronModelBridge): + + return provider + ++ def build_conversion_tasks(self, hf_pretrained, megatron_model): ++ """Inject virtual .weight keys so INT4 checkpoints (weight_packed/weight_scale/ ++ weight_zero_point) pass the hf_keys validation in the base class. ++ ++ When hf_checkpoint points to an INT4 compressed-tensors checkpoint, expert ++ weights are stored as weight_packed/weight_scale/weight_zero_point triplets ++ with no plain .weight key. The base build_conversion_tasks checks that each ++ mapped HF name exists in hf_keys and skips the param if not found, causing ++ all expert weights to be silently dropped. We patch get_all_keys() to return ++ synthetic .weight keys alongside the real packed keys so the check passes. ++ Downstream quantize_params in HfWeightIteratorBridge then converts the BF16 ++ output back to INT4 before sending to the rollout engine. ++ """ ++ if not (hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")): ++ return super().build_conversion_tasks(hf_pretrained, megatron_model) ++ ++ original_get_all_keys = hf_pretrained.state.source.get_all_keys ++ ++ def _get_all_keys_with_virtual(): ++ keys = original_get_all_keys() ++ all_keys_set = set(keys) ++ virtual_keys = [ ++ key[:-7] # "...weight_packed" -> "...weight" ++ for key in keys ++ if key.endswith("_packed") and f"{key[:-7]}_scale" in all_keys_set ++ ] ++ return keys + virtual_keys ++ ++ hf_pretrained.state.source.get_all_keys = _get_all_keys_with_virtual ++ try: ++ return super().build_conversion_tasks(hf_pretrained, megatron_model) ++ finally: ++ hf_pretrained.state.source.get_all_keys = original_get_all_keys ++ + def mapping_registry(self) -> MegatronMappingRegistry: + # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format + # First create simple 1:1 parameter mappings using a dictionary for readability +diff --git a/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +--- a/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py ++++ b/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +@@ -158,6 +158,7 @@ + *, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[Tensor] = None, ++ mtp_kwargs: Optional[dict] = None, + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, +@@ -259,6 +260,7 @@ + inference_context=inference_context, + output_processor=output_processor, + output_processor_context=output_processor_context, ++ mtp_kwargs=mtp_kwargs, + ) + + if _shadow_embedding: +diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py +index 8d65e29..e4af4f2 100644 +--- a/megatron/core/dist_checkpointing/strategies/torch.py ++++ b/megatron/core/dist_checkpointing/strategies/torch.py +@@ -501,10 +501,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner): + def _validate_global_shapes(self, metadata, sharded_tensors): + for sh_ten in sharded_tensors: + if sh_ten.key not in metadata.state_dict_metadata: +- raise KeyError( +- f"{sh_ten.key} from model not in state dict:" +- f" {sorted(metadata.state_dict_metadata.keys())}" +- ) ++ # raise KeyError( ++ # f"{sh_ten.key} from model not in state dict:" ++ # f" {sorted(metadata.state_dict_metadata.keys())}" ++ # ) ++ print(f"{sh_ten.key} from model not in state dict, will skip") ++ continue + loaded_shape = metadata.state_dict_metadata[sh_ten.key].size + expected_shape = sh_ten.global_shape + if loaded_shape != expected_shape: +@@ -528,7 +530,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner): + tensor_metadata = self.metadata.state_dict_metadata + metadata_with_sizes = [ + (tensor_metadata[key], tensor_metadata[key].size, sharded_tensor) +- for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() ++ for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() if key in tensor_metadata + ] + try: + # Temporarily set sizes to expected shapes +@@ -898,6 +900,7 @@ class TorchDistLoadShardedStrategy: + planner=MCoreLoadPlanner( + shapes_validation_sharded_tensors=flexible_shape_sharded_tensors, + allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors, ++ allow_partial_load=True, + flatten_state_dict=False, + flatten_sharded_tensors=False, + ), +diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py +index 8d797e8..98f6241 100644 +--- a/megatron/core/extensions/transformer_engine.py ++++ b/megatron/core/extensions/transformer_engine.py +@@ -918,6 +918,7 @@ class TELinear(te.pytorch.Linear): + ) + + for param in self.parameters(): ++ setattr(param, "parallel_mode", parallel_mode) + if is_expert: + # Reduce the gradient on the expert_data_parallel group for expert linear layers + setattr(param, "allreduce", not self.expert_parallel) +@@ -1948,6 +1949,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention): + + + if HAVE_TE and is_te_min_version("1.9.0.dev0"): ++ def ceil_div(x: int, y: int) -> int: ++ return (x + y - 1) // y ++ ++ class _FakeInt4QuantizationSTE(torch.autograd.Function): ++ @staticmethod ++ def forward(ctx, x, group_size): ++ m, n = x.shape ++ block_size_m, block_size_n = 1, group_size ++ ++ ++ m_padded = ceil_div(m, block_size_m) * block_size_m ++ n_padded = ceil_div(n, block_size_n) * block_size_n ++ ++ x_padded = torch.zeros( ++ (m_padded, n_padded), ++ dtype=x.dtype, device=x.device ++ ) ++ x_padded[:m, :n] = x ++ ++ x_view = x_padded.view( ++ m_padded // block_size_m, ++ block_size_m, ++ n_padded // block_size_n, ++ block_size_n ++ ) ++ ++ x_max = x_view.abs().float().amax(dim=(1, 3), keepdim=True) ++ q_max = 7 ++ x_scale = x_max / q_max ++ ++ x_scale = x_scale.clamp(min=1e-5) ++ ++ x_div = x_view / x_scale ++ x_round = torch.round(x_div) ++ ++ x_q_clamped = x_round.clamp(-q_max, q_max) ++ ++ x_dequant_view = x_q_clamped * x_scale ++ ++ x_dequant_full = x_dequant_view.view_as(x_padded) ++ x_out = x_dequant_full[:m, :n].contiguous().to(x.dtype) ++ ++ return x_out ++ ++ @staticmethod ++ def backward(ctx, grad_output): ++ return grad_output, None ++ ++ def fake_int4_quantization_ste(x, group_size): ++ x_out = _FakeInt4QuantizationSTE.apply(x, group_size) ++ ++ if hasattr(x, 'main_grad'): ++ x_out.main_grad = x.main_grad ++ ++ return x_out + + class TEGroupedLinear(te.pytorch.GroupedLinear): + """ +@@ -2179,6 +2235,7 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): + "amax_history_bwd": torch.cat( + [state["amax_history_bwd"].view(-1, 1) for state in state_list], + dim=1, ++ + ).view(self.fp8_meta["recipe"].amax_history_len, -1), + } + ) +@@ -2304,10 +2361,31 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): + def _get_weight_tensors(self): + """Get the weight tensors of the module.""" + weight_tensors = super()._get_weight_tensors() +- return maybe_fake_quantize_int4_weight_tensors( ++ ++ # Two independent int4 fake-QAT paths coexist here. Upstream's is ++ # config-driven and ships with the .dev.commit core; Relax's is the ++ # env-var one below, which is load-bearing -- documented in ++ # docs/zh/examples/low-precision-training.md and switched on by ++ # scripts/training/{text,multimodal}/run-{qwen3-30B-A3B,kimi-k2.6}*int4*.sh. ++ # The Relax patch appends its version verbatim, which alongside ++ # upstream's would define this method twice (second silently wins, ++ # disabling upstream's path). Compose them: each is a no-op unless its ++ # own switch is set. Enabling both at once would quantize twice -- a ++ # configuration error, not a case to arbitrate here. ++ weight_tensors = maybe_fake_quantize_int4_weight_tensors( + self.config, self.delay_wgrad_compute, weight_tensors + ) + ++ if os.getenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") == "1": ++ group_size = int(os.getenv("OPEN_TRAINING_INT4_GROUP_SIZE", "128")) ++ ++ weight_tensors = [ ++ fake_int4_quantization_ste(w, group_size) ++ for w in weight_tensors ++ ] ++ ++ return weight_tensors ++ + def _encode_extra_state(self, state): + # TE 2.0 changed the format of extra_state to be a byte tensor + if is_te_min_version("2.0.0"): +diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py +index d4cdaa2..31c7c91 100644 +--- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py ++++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py +@@ -513,6 +513,7 @@ def _mla_rope_fwd_kv_split_kernel( + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, ++ k_dim_ceil: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, +@@ -562,21 +563,27 @@ def _mla_rope_fwd_kv_split_kernel( + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + +- KV_ptr = KV + pid_m * stride_kv_seq + pid_head * BLOCK_H * stride_kv_nheads +- kv_off = tl.arange(0, BLOCK_H)[:, None] * stride_kv_nheads +- mask = kv_off < head_num * stride_kv_nheads +- k_in_off = kv_off + tl.arange(0, k_dim)[None, :] +- v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] +- k = tl.load(KV_ptr + k_in_off, mask=mask) +- v = tl.load(KV_ptr + v_in_off, mask=mask) ++ KV_ptr = KV + pid_m * stride_kv_seq # + pid_head * BLOCK_H * stride_kv_nheads ++ ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ kj_range = tl.arange(0, k_dim_ceil)[None, :] ++ mask_k = (ki_range < head_num) & (kj_range < k_dim) ++ mask_v = ki_range < head_num ++ k_off = ki_range * stride_kv_nheads + kj_range ++ if v_dim > 0: ++ v_off = ki_range * stride_kv_nheads + k_dim + tl.arange(0, v_dim)[None, :] ++ v = tl.load(KV_ptr + v_off, mask=mask_v) ++ else: ++ v = tl.zeros((BLOCK_H, 1), dtype=KV.dtype.element_ty) ++ k = tl.load(KV_ptr + k_off, mask=mask_k) + +- K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads +- V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads ++ K_ptr = O_KEY + pid_m * stride_k_seq # + pid_head * BLOCK_H * stride_k_nheads ++ V_ptr = O_VALUE + pid_m * stride_v_seq # + pid_head * BLOCK_H * stride_v_nheads + +- k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] +- v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] +- tl.store(K_ptr + k_out_off, k, mask=mask) +- tl.store(V_ptr + v_out_off, v, mask=mask) ++ k_out_off = ki_range * stride_k_nheads + kj_range ++ tl.store(K_ptr + k_out_off, k, mask=mask_k) ++ if v_dim > 0: ++ v_out_off = ki_range * stride_v_nheads + tl.arange(0, v_dim)[None, :] ++ tl.store(V_ptr + v_out_off, v, mask=mask_v) + + EMB = K_POS_EMB + pid_m * stride_emb_seq + # x1 = t[..., 0::2], x2 = t[..., 1::2] +@@ -588,24 +595,26 @@ def _mla_rope_fwd_kv_split_kernel( + x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + ++ x_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ mask_x = x_range < head_num + if REMOVE_INTERLEAVING: + x_1_off = ( +- tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads ++ x_range * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] * 2 + ) + x_2_off = x_1_off + 1 +- tl.store(K_ptr + x_1_off, x_left, mask=mask) +- tl.store(K_ptr + x_2_off, x_right, mask=mask) ++ tl.store(K_ptr + x_1_off, x_left, mask=mask_x) ++ tl.store(K_ptr + x_2_off, x_right, mask=mask_x) + else: + x_left_off = ( +- tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads ++ x_range * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 +- tl.store(K_ptr + x_left_off, x_left, mask=mask) +- tl.store(K_ptr + x_right_off, x_right, mask=mask) ++ tl.store(K_ptr + x_left_off, x_left, mask=mask_x) ++ tl.store(K_ptr + x_right_off, x_right, mask=mask_x) + + + @triton.autotune( +@@ -631,6 +640,7 @@ def _mla_rope_bwd_kv_split_kernel( + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, ++ k_dim_ceil: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, +@@ -672,27 +682,32 @@ def _mla_rope_bwd_kv_split_kernel( + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + +- dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads +- dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads +- mask = dkv_off < head_num * stride_dkv_nheads +- dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] +- dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] +- +- dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads +- dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads +- dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] +- dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] +- dk = tl.load(dK_ptr + dk_in_off, mask=mask) +- dv = tl.load(dV_ptr + dv_in_off, mask=mask) +- tl.store(dKV_ptr + dk_out_off, dk, mask=mask) +- tl.store(dKV_ptr + dv_out_off, dv, mask=mask) ++ dKV_ptr = dKV + pid_m * stride_dkv_seq # + pid_head * BLOCK_H * stride_dkv_nheads ++ ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ kj_range = tl.arange(0, k_dim_ceil)[None, :] ++ mask_k = (ki_range < head_num) & (kj_range < k_dim) ++ mask_v = ki_range < head_num ++ dk_out_off = ki_range * stride_dkv_nheads + kj_range ++ ++ dK_ptr = dK + pid_m * stride_dk_seq # + pid_head * BLOCK_H * stride_dk_nheads ++ dV_ptr = dV + pid_m * stride_dv_seq # + pid_head * BLOCK_H * stride_dv_nheads ++ dk_in_off = ki_range * stride_dk_nheads + kj_range ++ ++ dk = tl.load(dK_ptr + dk_in_off, mask=mask_k) ++ tl.store(dKV_ptr + dk_out_off, dk, mask=mask_k) ++ ++ if v_dim > 0: ++ dv_out_off = ki_range * stride_dkv_nheads + k_dim + tl.arange(0, v_dim)[None, :] ++ dv_in_off = ki_range * stride_dv_nheads + tl.arange(0, v_dim)[None, :] ++ dv = tl.load(dV_ptr + dv_in_off, mask=mask_v) ++ tl.store(dKV_ptr + dv_out_off, dv, mask=mask_v) + + if pid_head == 0: + x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): +- dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads +- x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim ++ dK_ptr = dK + pid_m * stride_dk_seq # + i * BLOCK_H * stride_dk_nheads ++ x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim + i * BLOCK_H * stride_dk_nheads + mask = x_off < head_num * stride_dk_nheads + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 +@@ -779,6 +794,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + + o_key = kv.new_empty(total_seqlen, nheads, emb_dim + k_dim) + o_value = kv.new_empty(total_seqlen, nheads, v_dim) ++ k_dim_ceil = triton.next_power_of_2(k_dim) + + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) + _mla_rope_fwd_kv_split_kernel[grid]( +@@ -790,6 +806,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + sin, + emb_dim, + k_dim, ++ k_dim_ceil, + v_dim, + nheads, + batch_size, +@@ -849,6 +866,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + + d_kv = dk.new_empty(total_seqlen, nheads, ctx.k_dim + ctx.v_dim) + d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) ++ k_dim_ceil = triton.next_power_of_2(ctx.k_dim) + + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) + _mla_rope_bwd_kv_split_kernel[grid]( +@@ -860,6 +878,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + sin, + ctx.emb_dim, + ctx.k_dim, ++ k_dim_ceil, + ctx.v_dim, + nheads, + batch_size, +diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py +index 4b3e360..c31c80b 100644 +--- a/megatron/core/inference/contexts/dynamic_context.py ++++ b/megatron/core/inference/contexts/dynamic_context.py +@@ -67,7 +67,8 @@ except ImportError: + try: + from torch_memory_saver import torch_memory_saver + +- torch_memory_saver.hook_mode = "torch" ++ # Commented out: breaks SGLang CUDA graph (requires hook_mode="preload") ++ # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = True + except ImportError: + HAVE_TORCH_MEMORY_SAVER = False +diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py +index 5b938cc..b681685 100755 +--- a/megatron/core/models/gpt/gpt_layer_specs.py ++++ b/megatron/core/models/gpt/gpt_layer_specs.py +@@ -195,6 +195,8 @@ def get_gpt_layer_with_transformer_engine_submodules( + mla_down_proj_fusion: bool = False, + dense_grouped_gemm: bool = False, + use_grouped_gemm_for_dense_mlp: bool = False, ++ post_self_attn_layernorm: bool = False, ++ post_mlp_layernorm: bool = False, + ) -> TransformerLayerSubmodules: + """Use these submodules to use lower-level Transformer Engine modules (required for fp8 + training). +@@ -289,9 +291,11 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + ), + self_attn_bda=get_bias_dropout_add, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + sharded_state_dict_keys_map=( + { + "self_attention.linear_q_down_proj.layer_norm_": "input_layernorm.", +@@ -321,10 +325,12 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm(has_residual=True) if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + ) + else: + qk_norm = backend.layer_norm(for_qk=True) +@@ -346,10 +352,12 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm(has_residual=True) if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + sharded_state_dict_keys_map={ + "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight", + "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias", +diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py +--- a/megatron/core/models/gpt/gpt_model.py ++++ b/megatron/core/models/gpt/gpt_model.py +@@ -34,6 +34,7 @@ + MultiTokenPredictionBlock, + mtp_on_this_rank, + process_mtp_loss, ++ roll_tensor, + ) + from megatron.core.transformer.spec_utils import ModuleSpec + from megatron.core.transformer.transformer_block import TransformerBlock +@@ -529,6 +530,7 @@ + padding_mask: Optional[Tensor] = None, + output_processor: Optional[Callable[..., Tensor]] = None, + output_processor_context: Optional[Any] = None, ++ mtp_kwargs: Optional[dict] = None, + ) -> Tensor: + """Forward function of the GPT Model This function passes the input tensors + through the embedding layer, and then the decoder and finally into the post +@@ -624,6 +626,7 @@ + mhc_multistream=mhc_multistream, + output_processor=output_processor, + output_processor_context=output_processor_context, ++ mtp_kwargs=mtp_kwargs, + ) + + def _postprocess( +@@ -649,6 +652,7 @@ + mhc_multistream=None, + output_processor=None, + output_processor_context=None, ++ mtp_kwargs=None, + ): + """Postprocesses decoder hidden states to generate logits or compute loss. + +@@ -673,7 +677,42 @@ + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() +- if mtp_in_postprocess and not (in_inference_mode or is_spec_decode): ++ # Relax always calls forward() with labels=None and computes the main loss ++ # itself, so it hands MTP its labels explicitly via `mtp_kwargs`. Upstream's ++ # `process_mtp_loss(input_ids=...)` fallback is not a substitute: on the ++ # VL+THD+CP unsplit path `input_ids` is the un-CP-split `unsplit_tokens` ++ # (twice the per-chunk length), so the derived labels do not line up with ++ # the CP-split hidden states and vocab_parallel_cross_entropy raises ++ # "shape mismatch: indexing tensors could not be broadcast together". ++ mtp_labels = labels ++ if mtp_kwargs is not None and mtp_kwargs.get("mtp_labels", None) is not None: ++ mtp_labels = mtp_kwargs["mtp_labels"] ++ mtp_labels, _ = roll_tensor( ++ mtp_labels, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.pg_collection.cp, ++ packed_seq_params=packed_seq_params, ++ ) ++ if loss_mask is not None: ++ loss_mask, _ = roll_tensor( ++ loss_mask, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.pg_collection.cp, ++ packed_seq_params=packed_seq_params, ++ ) ++ # `mtp_labels is not None` gates BOTH the MTP block below and the ++ # process_mtp_loss() call further down: Relax's log-prob forward ++ # (model.py forward_step) passes labels=None and no mtp_kwargs, and the ++ # old patch relied on process_mtp_loss returning early in that case. ++ # Upstream replaced that early return with the input_ids fallback, so the ++ # skip has to happen here instead. ++ if ( ++ mtp_in_postprocess ++ and not (in_inference_mode or is_spec_decode) ++ and mtp_labels is not None ++ ): + hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, +@@ -694,7 +733,7 @@ + if not self.post_process: + return hidden_states + +- if self.config.mtp_num_layers: ++ if self.config.mtp_num_layers and mtp_labels is not None: + assert self.config.mtp_num_layers > 0 + if in_inference_mode or is_spec_decode: + # Cache decoder hidden states for serial MTP computation +@@ -705,7 +744,7 @@ + mtp_cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params) + hidden_states = process_mtp_loss( + hidden_states=hidden_states, +- labels=labels, ++ labels=mtp_labels, + loss_mask=loss_mask, + output_layer=self.output_layer, + output_weight=output_weight, +diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +--- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py ++++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +@@ -249,6 +249,11 @@ + return cpu_optimizers + + def _get_sub_optimizer_param_groups(self, offload_fraction: float): ++ import warnings ++ warnings.warn( ++ "CPU offload optimizer init can be very slow (potentially minutes) for " ++ "large MoE models due to per-parameter pinned-memory allocation and H2D copies." ++ ) + params = [] + for group in self.param_groups: + params.extend(group["params"]) +@@ -370,8 +375,11 @@ + if not self.param_update_in_fp32: + return + for param, v in self.state.items(): +- fp32_param = self.param_to_fp32_param[param] +- fp32_param.data.copy_(v["master_param"]) ++ # Native FP32 params do not need a separate master parameter and are ++ # intentionally absent from param_to_fp32_param. ++ fp32_param = self.param_to_fp32_param.get(param) ++ if fp32_param is not None: ++ fp32_param.data.copy_(v["master_param"]) + + def update_fp32_param_by_new_param(self): + """ +@@ -379,6 +387,16 @@ + """ + for param, fp32_param in self.param_to_fp32_param.items(): + fp32_param.data.copy_(param) ++ # Params that are already fp32 never get an entry in param_to_fp32_param ++ # (see _get_sub_optimizer_param_groups), so the loop above misses their CPU ++ # snapshot. That snapshot is taken at construction time and copied back over ++ # the model param by the post-step hook, which silently reverts anything ++ # written afterwards -- a checkpoint load, most obviously -- to the values the ++ # model was initialised with. For non-fp32 params the two maps share one ++ # tensor, so this loop only needs to cover the rest. ++ for param, cpu_copy in self.gpu_params_map_cpu_copy.items(): ++ if param not in self.param_to_fp32_param: ++ cpu_copy.data.copy_(param) + + def _register_load_state_dict_hooks(self): + def pre_load_state_dict_hook(self, state_dict): +diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py +index 9baed75..8f68e3d 100644 +--- a/megatron/core/optimizer/distrib_optimizer.py ++++ b/megatron/core/optimizer/distrib_optimizer.py +@@ -471,7 +471,9 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + + # fp32 params. + elif model_param.type() == 'torch.cuda.FloatTensor': +- shard_model_param = model_param.view(-1)[param_range.start : param_range.end] ++ shard_model_param = model_param.detach().view(-1)[ ++ param_range.start : param_range.end ++ ] + model_fp32_params_this_group.append(model_param) + shard_fp32_params_this_group.append(shard_model_param) + tensor_parallel.copy_tensor_model_parallel_attributes( +@@ -888,6 +890,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + # TE FusedAdam will not accumulate step for empty param groups, so we need to + # align the step across param groups. + param_group["step"] = int(step) ++ if "step" in param_group and param_group["step"] is None: ++ del param_group["step"] + + # Grad scaler state. + if self.grad_scaler: +@@ -1971,6 +1975,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + # separately via param_groups, not as part of the gradient buffer. + tensors[key] = LocalNonpersistentObject(tensors[key]) + continue ++ if key == 'step': ++ continue + assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), ( + tensors[key].shape, + gbuf_local_start, +diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py +index 7023488..25b07e2 100644 +--- a/megatron/core/parallel_state.py ++++ b/megatron/core/parallel_state.py +@@ -11,6 +11,7 @@ from typing import Callable, List, Optional + + import numpy as np + import torch ++import torch.distributed as dist + + from megatron.core.inference.symmetric_memory import SymmetricMemoryManager + +diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py +index 465e83f..232caef 100644 +--- a/megatron/core/pipeline_parallel/p2p_communication.py ++++ b/megatron/core/pipeline_parallel/p2p_communication.py +@@ -232,34 +232,26 @@ class P2PCommunicator: + group=self.pp_group, + ) + else: +- ops = [] +- if send_prev_shape_tensor is not None: +- send_prev_op = torch.distributed.P2POp( +- torch.distributed.isend, send_prev_shape_tensor, self.prev_rank, self.pp_group +- ) +- ops.append(send_prev_op) +- if recv_prev_shape_tensor is not None: +- recv_prev_op = torch.distributed.P2POp( +- torch.distributed.irecv, recv_prev_shape_tensor, self.prev_rank, self.pp_group +- ) +- ops.append(recv_prev_op) +- if send_next_shape_tensor is not None: +- send_next_op = torch.distributed.P2POp( +- torch.distributed.isend, send_next_shape_tensor, self.next_rank, self.pp_group +- ) +- ops.append(send_next_op) +- if recv_next_shape_tensor is not None: +- recv_next_op = torch.distributed.P2POp( +- torch.distributed.irecv, recv_next_shape_tensor, self.next_rank, self.pp_group +- ) +- ops.append(recv_next_op) +- if len(ops) > 0: +- reqs = torch.distributed.batch_isend_irecv(ops) +- for req in reqs: +- req.wait() ++ # PR #5271 (Megatron-LM): shape exchange MUST use _p2p_ops rather than ++ # batch_isend_irecv. batch_isend_irecv is one tagless NCCL group; when ++ # pp_group.size()==2, prev_rank == next_rank (single physical peer) and ++ # same-peer ops pair FIFO by enqueue order. Both ranks build ops in the ++ # same fixed order but hold opposite prev/next roles → recv_prev_shape ++ # and recv_next_shape get silently crossed. _p2p_ops handles size==2 ++ # via the group.WORLD split + even/odd ordering, correct for size>=4 too. ++ reqs = _p2p_ops( ++ tensor_send_prev=send_prev_shape_tensor, ++ tensor_recv_prev=recv_prev_shape_tensor, ++ tensor_send_next=send_next_shape_tensor, ++ tensor_recv_next=recv_next_shape_tensor, ++ group=self.pp_group, ++ prev_pipeline_rank=self.prev_rank, ++ next_pipeline_rank=self.next_rank, ++ ) ++ for req in reqs.values(): ++ req.wait() + +- # To protect against race condition when using batch_isend_irecv(). +- # should take this out once the bug with batch_isend_irecv is resolved. ++ # keep the CUDA sync as a defensive measure — cheap for 3-int64 tensors. + torch.cuda.synchronize() + + recv_prev_shape = [0, 0, 0] +@@ -371,6 +363,11 @@ class P2PCommunicator: + return [] + + p2p_func = _ring_exchange_wrapper ++ elif self.pp_group.size() == 2: ++ # PR #5271 (Megatron-LM): size==2 has same-peer (prev_rank == next_rank); ++ # batch_isend_irecv cannot pair same-peer bidirectional ops correctly (see ++ # #1450). Force _p2p_ops which handles this via WORLD split + even/odd order. ++ p2p_func = _p2p_ops + elif config.batch_p2p_comm: + assert wait_on_reqs + p2p_func = _batched_p2p_ops +@@ -381,10 +378,12 @@ class P2PCommunicator: + next_rank = self.next_rank + prev_rank = self.prev_rank + +- if config.use_ring_exchange_p2p or config.batch_p2p_comm: +- reqs = [] +- else: ++ # reqs init must match p2p_func return type: _p2p_ops returns dict, ++ # _batched_p2p_ops + _ring_exchange_wrapper return list. ++ if p2p_func is _p2p_ops: + reqs = {} ++ else: ++ reqs = [] + + tensor_recv_prev = None + tensor_recv_next = None +diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py +index 7e28691..3c19562 100644 +--- a/megatron/core/ssm/gated_delta_net.py ++++ b/megatron/core/ssm/gated_delta_net.py +@@ -584,16 +584,17 @@ class GatedDeltaNet(MegatronModule): + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) +- query, key, value, gate, beta, g = self.pre_gated_delta_rule( +- qkvzba, +- batch, +- seq_len_post_headwise, +- cp_size_headwise, +- cp_group_headwise, +- cu_seqlens_q, +- chunkwise_cp_context, +- packed_seq_params=packed_seq_params, +- ) ++ with torch._dynamo.config.patch(disable=True): ++ query, key, value, gate, beta, g = self.pre_gated_delta_rule( ++ qkvzba, ++ batch, ++ seq_len_post_headwise, ++ cp_size_headwise, ++ cp_group_headwise, ++ cu_seqlens_q, ++ chunkwise_cp_context, ++ packed_seq_params=packed_seq_params, ++ ) + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") +@@ -1211,7 +1212,7 @@ def get_parameter_local_cp_headwise( + slices = [slice(None)] * param.dim() + dim_size = param.size(dim=dim) + slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) +- param = param[slices] ++ param = param[tuple(slices)] + return param + + +diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py +index dbecd73..1c52478 100644 +--- a/megatron/core/tensor_parallel/random.py ++++ b/megatron/core/tensor_parallel/random.py +@@ -589,7 +589,30 @@ class CheckpointFunction(torch.autograd.Function): + ) + + # Store everything. +- ctx.save_for_backward(*args) ++ # ++ # save_for_backward accepts tensors only, but not every checkpointed ++ # argument is one: gemma-4 uses dual RoPE, so TransformerBlock passes ++ # rotary_pos_emb as a TUPLE of two tensors, and this raised ++ # TypeError: save_for_backward can only save variables, ++ # but argument 4 is of type tuple ++ # Split the non-tensor arguments out here and rebuild the original ++ # ordering in backward(). ++ # ++ # Deliberately NOT calling _save_args_to_ctx() defined below, even ++ # though it solves the same problem for CheckpointWithoutOutputFunction: ++ # that helper detaches at save time, whereas this Function detaches in ++ # backward via detach_variable(). Splitting inline keeps the all-tensor ++ # path byte-identical to upstream, so no other model changes behaviour. ++ tensor_args = [] ++ non_tensor_entries = [] ++ for index, arg in enumerate(args): ++ if isinstance(arg, torch.Tensor): ++ tensor_args.append(arg) ++ else: ++ non_tensor_entries.append((index, arg)) ++ ctx.save_for_backward(*tensor_args) ++ ctx.non_tensor_entries = tuple(non_tensor_entries) ++ ctx.total_args_count = len(args) + + _unset_checkpointing() + return outputs +@@ -607,7 +630,17 @@ class CheckpointFunction(torch.autograd.Function): + ) + _set_checkpointing() + +- inputs = ctx.saved_tensors ++ # Rebuild the forward arguments in their original order, re-inserting the ++ # non-tensor ones that forward() kept out of save_for_backward. When every ++ # argument was a tensor this is exactly ctx.saved_tensors. ++ # detach_variable() (torch.utils.checkpoint) passes non-tensors through ++ # untouched, so the recompute below needs no further changes. ++ _saved = iter(ctx.saved_tensors) ++ _non_tensors = dict(ctx.non_tensor_entries) ++ inputs = tuple( ++ _non_tensors[i] if i in _non_tensors else next(_saved) ++ for i in range(ctx.total_args_count) ++ ) + if ctx.distribute_saved_activations: + safely_set_viewless_tensor_data( + inputs[0], gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape) +@@ -630,7 +663,11 @@ class CheckpointFunction(torch.autograd.Function): + *filter(lambda x: torch.is_tensor(x[0]) and x[0].requires_grad, zip(outputs, args)) + ) + torch.autograd.backward(outputs, args) +- grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs) ++ # One entry per forward argument. autograd requires None -- not the object ++ # itself -- for arguments that are not tensors; returning `inp` here would ++ # raise "expected Variable or None". Unreachable before this commit, ++ # because a non-tensor argument died in save_for_backward first. ++ grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in detached_inputs) + + _unset_checkpointing() + return (None, None) + grads +diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +index 3459a19..1a4f130 100644 +--- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py ++++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +@@ -276,4 +276,3 @@ def prepare_cp_compressor_input( +-@torch.compile + def _build_cp_indexer_layout( + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index c8e197d..c5928f1 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -797,6 +797,9 @@ def topk_routing_with_score_function( + scores, topk, num_groups, group_topk, _compute_topk + ) + ++ from relax.utils.training.routing_replay import get_routing_replay_compute_topk ++ compute_topk = get_routing_replay_compute_topk(compute_topk) ++ + # Precision notes: + # - Logits are converted to fp32 for score functions. + # - All the intermediate calculations are in fp32. +@@ -1356,7 +1359,11 @@ class RouterGatingLinearFunction(torch.autograd.Function): + inp_shape = inp.shape + inp = inp.view(-1, inp_shape[-1]) + +- if te_general_gemm is not None and router_dtype != torch.float64: ++ # TE multiplies in the operand dtype, so BF16 operands silently defeat ++ # --moe-router-dtype fp32. Keep FP32/FP64 routing on torch.mm. ++ if te_general_gemm is not None and router_dtype not in ( ++ torch.float32, torch.float64 ++ ): + output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias) + output = output[0] + elif bias is None: +@@ -1389,7 +1396,9 @@ class RouterGatingLinearFunction(torch.autograd.Function): + inp = inp.view(-1, inp_shape[-1]) + grad_output = grad_output.view(-1, grad_shape[-1]) + +- if te_general_gemm is not None and ctx.router_dtype != torch.float64: ++ if te_general_gemm is not None and ctx.router_dtype not in ( ++ torch.float32, torch.float64 ++ ): + grad_input = te_general_gemm( + weight.to(ctx.router_dtype), grad_output, ctx.router_dtype, layout="NN", grad=True + ) +diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py +index 2796bc6..71e036e 100644 +--- a/megatron/core/transformer/moe/router.py ++++ b/megatron/core/transformer/moe/router.py +@@ -267,6 +267,9 @@ class TopKRouter(Router): + if self.config.moe_enable_routing_replay: + self.router_replay = RouterReplay() + ++ from relax.utils.training.routing_replay import register_routing_replay ++ register_routing_replay(self) ++ + def _maintain_float32_expert_bias(self): + """ + Maintain the expert bias in float32. +diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py +--- a/megatron/core/transformer/multi_token_prediction.py ++++ b/megatron/core/transformer/multi_token_prediction.py +@@ -1031,7 +1031,7 @@ + ) + derived_labels_from_input_ids = True + +- if config.mtp_detach_heads: ++ if getattr(config, "mtp_detach_lm_head", True): + if output_weight is not None: + output_weight = output_weight.detach() + else: +@@ -1084,6 +1084,7 @@ + mtp_logits = scale_logits_fn(mtp_logits) + mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits) + mtp_loss = loss_mask * mtp_loss ++ mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers + + if is_training: + if mtp_logits is not None: +@@ -1097,7 +1098,7 @@ + total = torch.zeros((), device=mtp_loss.device, dtype=mtp_loss.dtype) + + MTPLossLoggingHelper.save_loss_to_tracker( +- torch.sum(mtp_loss), ++ mtp_loss_scale * torch.sum(mtp_loss), + num_tokens, + mtp_layer_number, + config.mtp_num_layers, +@@ -1106,7 +1107,6 @@ + avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True), + calculate_per_token_loss=config.calculate_per_token_loss, + ) +- mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers + if config.calculate_per_token_loss: + # When calculate_per_token_loss is enabled, finalize_model_grads will + # divide all gradients by total_num_tokens (from main loss). +@@ -1383,12 +1383,12 @@ + # embedding + decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) + +- if self.config.mtp_detach_heads: ++ if getattr(self.config, "mtp_detach_embedding", True): + decoder_input = decoder_input.detach() + + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + # make_viewless_tensor no-ops when hidden_states is not a view (_base is None), +- # which happens after detach() with mtp_detach_heads. Activation ++ # which happens after detach() with mtp_detach_embedding. Activation + # checkpointing (CheckpointFunction.apply) requires at least one input tensor + # with requires_grad=True to produce a differentiable output, so we ensure it + # here to maintain gradient flow to MTP layer parameters. +@@ -2167,8 +2167,11 @@ + else: + hidden_states = hidden_states_list[offset] + +- if self.config.mtp_detach_heads: +- hidden_states = hidden_states.detach() ++ if offset == 0 and getattr(self.config, "mtp_detach_backbone", True): ++ # offset > 0 means hidden_states came from a previous MTP stage under VPP, ++ # not from the main backbone -- detaching there would sever gradient flow ++ # between MTP layers. ++ hidden_states = hidden_states.detach().requires_grad_(True) + + for iteration in range(self.config.mtp_num_layers): + layer_idx = 0 if self.mtp_use_repeated_layer else iteration +diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py +index 43e45b7..8166f94 100644 +--- a/megatron/core/transformer/transformer_config.py ++++ b/megatron/core/transformer/transformer_config.py +@@ -224,6 +224,10 @@ class TransformerConfig(ModelParallelConfig): + """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func + is quick_gelu or weighted SwiGLU (MoE only).""" + ++ activation_func_clamp_shared_expert: bool = True ++ """If False, skip activation_func_clamp_value inside SharedExpertMLP so only routed MoE ++ experts get the clamp.""" ++ + 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.""" +@@ -283,6 +287,9 @@ class TransformerConfig(ModelParallelConfig): + num_query_groups >= tp, and under fp8/fp4 a per-partition + linear_qkv_out_dim aligned to 16/32.""" + ++ post_self_attn_layernorm: bool = False ++ post_mlp_layernorm: bool = False ++ + test_mode: bool = False + """Whether to run real-time tests.""" + +diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py +index 34a456b..8182c7c 100644 +--- a/megatron/core/transformer/transformer_layer.py ++++ b/megatron/core/transformer/transformer_layer.py +@@ -280,6 +280,7 @@ class TransformerLayerSubmodules: + self_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp + self_attention: Union[ModuleSpec, type] = IdentityOp + self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp ++ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp + + pre_cross_attn_layernorm: LayerNormBuilder = IdentityOp + cross_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp +@@ -290,6 +291,7 @@ class TransformerLayerSubmodules: + mlp_hyper_connection: Union[ModuleSpec, type] = IdentityOp + mlp: MlpBuilder | type[IdentityOp] = IdentityOp + mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp ++ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp + + # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method + sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict) +@@ -395,6 +397,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + # [Module 3: BiasDropoutFusion] + self.self_attn_bda = build_module(submodules.self_attn_bda) + ++ self.post_self_attn_layernorm = build_module( ++ submodules.post_self_attn_layernorm, ++ config=self.config, ++ hidden_size=self.config.hidden_size, ++ eps=self.config.layernorm_epsilon, ++ ) ++ + # [Module 4: Post SelfAttention] Optional Layernorm after self-attn + self.pre_cross_attn_layernorm = submodules.pre_cross_attn_layernorm( + config=self.config, +@@ -466,6 +475,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + + self.is_moe_layer = isinstance(self.mlp, MoELayer) + ++ self.post_mlp_layernorm = build_module( ++ submodules.post_mlp_layernorm, ++ config=self.config, ++ hidden_size=self.config.hidden_size, ++ eps=self.config.layernorm_epsilon ++ ) ++ + self.recompute_input_layernorm = False + self.recompute_pre_mlp_layernorm = False + self.recompute_mlp = False +@@ -780,6 +796,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + attention_output_with_bias[0] + ) + ++ attention_output, attention_output_bias = attention_output_with_bias ++ attention_output = self.post_self_attn_layernorm(attention_output) ++ attention_output_with_bias = (attention_output, attention_output_bias) ++ + # TODO: could we move `bias_dropout_add_exec_handler` itself + # inside the module provided in the `bias_dropout_add_spec` module? + nvtx_range_push(suffix="self_attn_bda") +@@ -1036,6 +1056,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + ) + mlp_output_with_bias = (mlp_output, mlp_bias) + ++ mlp_output, mlp_output_bias = mlp_output_with_bias ++ mlp_output = self.post_mlp_layernorm(mlp_output) ++ mlp_output_with_bias = (mlp_output, mlp_output_bias) ++ + nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias, residual + +diff --git a/megatron/training/training.py b/megatron/training/training.py +index af94ff6..b309bfa 100644 +--- a/megatron/training/training.py ++++ b/megatron/training/training.py +@@ -221,7 +221,9 @@ except ImportError: + try: + from torch_memory_saver import torch_memory_saver + +- torch_memory_saver.hook_mode = "torch" ++ # NOTE(wuhuan): keep the default hook mode; forcing "torch" triggers ++ # 'torch.AcceleratorError: CUDA error: invalid argument' on weight updates. ++ # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = True + except ImportError: + HAVE_TORCH_MEMORY_SAVER = False +diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py +--- a/megatron/core/transformer/moe/shared_experts.py ++++ b/megatron/core/transformer/moe/shared_experts.py +@@ -123,6 +123,9 @@ class SharedExpertMLP(MLP): + assert config.add_bias_linear == False, "bias is not supported in the shared experts, " + "please set '--disable-bias-linear' instead." + ++ if not config.activation_func_clamp_shared_expert: ++ config.activation_func_clamp_value = None ++ + config.ffn_hidden_size = config.moe_shared_expert_intermediate_size + # TODO(Hepteract): pass pg_collection to MLP after refactoring MLP + super().__init__(config=config, submodules=submodules, tp_group=pg_collection.tp, name=name) +diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py +--- a/megatron/core/fp8_utils.py ++++ b/megatron/core/fp8_utils.py +@@ -339,2 +339,11 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, + ) ++ ++ recipe_configurator = getattr(quantizer_factory, "configure_custom_recipe", None) ++ if recipe_configurator is not None: ++ if not callable(recipe_configurator): ++ raise ValueError( ++ "fp8 quantizer factory configure_custom_recipe attribute must be callable." ++ ) ++ recipe_configurator(custom_recipe) ++ + return custom_recipe +diff --git a/megatron/core/safe_globals.py b/megatron/core/safe_globals.py +--- a/megatron/core/safe_globals.py ++++ b/megatron/core/safe_globals.py +@@ -102,13 +102,15 @@ class SafeUnpickler(pickle.Unpickler): + ("torch.storage", "_load_from_bytes"), + ("transformer_engine.common.recipe", "DelayedScaling"), + ("transformer_engine.common.recipe", "Float8CurrentScaling"), + ("transformer_engine.common.recipe", "Float8BlockScaling"), ++ ("transformer_engine.common.recipe", "CustomRecipe"), + ("transformer_engine.common.recipe", "MXFP8BlockScaling"), + ("transformer_engine.common.recipe", "NVFP4BlockScaling"), + ("transformer_engine.common.recipe", "Format"), + ("transformer_engine.common.recipe", "_FormatHelper"), + ("transformer_engine.common.recipe", "MMParams"), + ("transformer_engine.common.recipe", "QParams"), ++ ("relax.backends.megatron.fp8_recipes", "_DeepSeekV4SenderAlignedQuantizerFactory"), + ("megatron.core.extensions.transformer_engine", "TEDelayedScaling"), + ("megatron.core.safe_globals", "safe_load_from_bytes"), + ("numpy._core.multiarray", "_reconstruct"), +diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py +--- a/megatron/core/transformer/experimental_attention_variant/csa.py ++++ b/megatron/core/transformer/experimental_attention_variant/csa.py +@@ -1054,2 +1054,4 @@ class Compressor(MegatronModule): +- kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) +- score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) ++ # Run the compressor GEMMs in high precision (BF16) even under FP8 training. ++ with get_fp8_disabled_context(self.config): ++ kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) ++ score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) +diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +--- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py ++++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +@@ -261,6 +261,30 @@ + return global_idxs.int() + + ++def _compact_flat_topk_idxs(global_idxs: Tensor) -> Tuple[Tensor, Tensor]: ++ """Pack valid global indices into a per-row prefix. ++ ++ The returned ``topk_length`` selects that prefix in FlashMLA forward and ++ cuDNN DSA backward. Invalid suffix entries remain ``-1`` until forward has ++ consumed them; callers may then replace the ignored suffix with a safe ++ non-negative placeholder before backward. ++ """ ++ if global_idxs.ndim != 2: ++ raise ValueError(f"global_idxs must be 2-D (rows, topk), got {tuple(global_idxs.shape)}") ++ ++ if global_idxs.is_cuda: ++ _ensure_dsa_namespace() ++ res = _DSA.compactify_wrapper(global_idxs) ++ compact_idxs, topk_length = res["indices"], res["topk_length"] ++ else: ++ valid_mask = global_idxs >= 0 ++ sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) ++ compact_idxs = global_idxs.gather(-1, sorted_indices) ++ topk_length = valid_mask.sum(dim=-1).int() ++ ++ return compact_idxs.int().contiguous(), topk_length.int().contiguous() ++ ++ + def build_flat_topk_idxs( + *idx_groups: Tensor, + batch_size: int, +@@ -309,21 +333,9 @@ + + topk_length_flat = None + if compact: +- if global_idxs.is_cuda: +- # Fast path: single warp-per-row CuTe DSL kernel from cuDNN's DSA +- # namespace. Replaces a stable argsort + gather + sum + permute +- # chain with one global-load + global-store per element. +- _ensure_dsa_namespace() +- res = _DSA.compactify_wrapper(global_idxs) +- global_idxs, topk_length_flat = res["indices"], res["topk_length"] +- else: +- # CPU fallback so the unit tests that exercise this helper without +- # CUDA still work. Production callers always go through the CUDA +- # path above. +- valid_mask = global_idxs >= 0 +- sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) +- global_idxs = global_idxs.gather(-1, sorted_indices) +- topk_length_flat = valid_mask.sum(dim=-1).int() ++ # The CUDA path is a single warp-per-row CuTe DSL kernel; the helper ++ # retains a stable PyTorch fallback for CPU-only unit tests. ++ global_idxs, topk_length_flat = _compact_flat_topk_idxs(global_idxs) + + return global_idxs, topk_length_flat + +@@ -1136,6 +1148,14 @@ + combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1) + global_idxs = local_to_global_flat(combined_local, b) + ++ # When the teacher/indexer loss is disabled, compact the complete ++ # attention set once and share its valid-prefix length between ++ # FlashMLA forward and cuDNN DSA backward. Keep the legacy segmented ++ # layout for loss_coeff > 0 until the full teacher-loss fix is ported. ++ topk_length = None ++ if loss_coeff == 0: ++ global_idxs, topk_length = _compact_flat_topk_idxs(global_idxs) ++ + # ---- 4. FlashMLA forward (flat layout for both SBHD and THD). -------- + if is_thd: + q_flat = query +@@ -1149,9 +1169,13 @@ + global_idxs, + softmax_scale, + attn_sink=attn_sink, +- topk_length=None, +- indexer_topk=indexer_topk, ++ topk_length=topk_length, ++ indexer_topk=0 if topk_length is not None else indexer_topk, + ) ++ if topk_length is not None: ++ # Forward has consumed the -1 sentinels. cuDNN DSA backward still ++ # reads ignored suffix slots, so replace them with a safe address. ++ global_idxs.clamp_min_(0) + + # ---- 4b. Derive padding-row mask for loss exclusion. ----------------- + # When CUDA-graph padding makes cu_seqlens_q cover all total_q rows +@@ -1173,6 +1197,52 @@ + real_len_per_row = real_seg_lens[row_batch_ids].to(torch.int32) + padding_row_mask = pos_in_seg >= real_len_per_row + ++ if topk_length is not None: ++ # cuDNN DSA backward requires at least one tile. The masked ++ # dO/LSE below makes this harmless placeholder gradient-free. ++ topk_length.masked_fill_(padding_row_mask, 1) ++ ++ if loss_coeff == 0: ++ # The old teacher path unconditionally consumes FlashMLA's partial ++ # indexer LSE, which is unavailable in compact mode. Since the ++ # configured coefficient is zero, skip only that inactive work; ++ # loss_coeff > 0 retains the legacy teacher behavior unchanged. ++ indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) ++ precomputed_grad_q_indexer = torch.zeros_like(q_indexer) ++ precomputed_grad_k_indexer = torch.zeros_like(k_indexer) ++ precomputed_grad_weights = torch.zeros_like(weights) ++ ctx.save_for_backward( ++ q_flat, ++ kv_flat, ++ attn_sink, ++ global_idxs, ++ topk_length, ++ out_flat, ++ lse, ++ precomputed_grad_q_indexer, ++ precomputed_grad_k_indexer, ++ precomputed_grad_weights, ++ ) ++ ctx.softmax_scale = softmax_scale ++ ctx.is_thd = is_thd ++ ctx.has_topk_length = True ++ ctx.padding_row_mask = padding_row_mask ++ ctx.np_ = np_ ++ ctx.d = d ++ if is_thd: ++ ctx.total_q = total_q ++ else: ++ ctx.sq = sq ++ ctx.b = b ++ ctx.skv = skv ++ ++ d_v = out_flat.shape[-1] ++ if is_thd: ++ output = out_flat.reshape(total_q, np_ * d_v) ++ else: ++ output = out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v) ++ return output, indexer_loss ++ + # ---- 5. Derive predict from indexer_scores, compute target. ---------- + # Layout-specific attn tensors (detached — loss is not differentiable + # through them). +@@ -1383,11 +1453,13 @@ + precomputed_grad_weights[padding_row_mask] = 0 + + # ---- 7. Save context (only sparse-attn bwd tensors + indexer grads). - ++ empty_topk_length = torch.empty(0, dtype=torch.int32, device=global_idxs.device) + ctx.save_for_backward( + q_flat, + kv_flat, + attn_sink, + global_idxs, ++ empty_topk_length, + out_flat, + lse, + precomputed_grad_q_indexer, +@@ -1396,6 +1468,8 @@ + ) + ctx.softmax_scale = softmax_scale + ctx.is_thd = is_thd ++ ctx.has_topk_length = False ++ ctx.padding_row_mask = padding_row_mask + ctx.np_ = np_ + ctx.d = d + if is_thd: +@@ -1421,12 +1495,14 @@ + kv_flat, + attn_sink, + global_idxs, ++ saved_topk_length, + out_flat, + lse, + precomputed_grad_q_indexer, + precomputed_grad_k_indexer, + precomputed_grad_weights, + ) = ctx.saved_tensors ++ topk_length = saved_topk_length if ctx.has_topk_length else None + + is_thd = ctx.is_thd + np_, d = ctx.np_, ctx.d +@@ -1439,6 +1515,10 @@ + sq, b, skv = ctx.sq, ctx.b, ctx.skv + dO_flat = grad_output.reshape(sq * b, np_, d_v) + ++ if topk_length is not None and ctx.padding_row_mask is not None: ++ dO_flat = dO_flat.masked_fill(ctx.padding_row_mask[:, None, None], 0) ++ lse = lse.masked_fill(ctx.padding_row_mask[:, None], 0) ++ + attn_bwd = _DSA.sparse_attention_backward_wrapper( + q_flat, + kv_flat, +@@ -1448,7 +1528,7 @@ + attn_sink, + global_idxs, + softmax_scale=ctx.softmax_scale, +- topk_length=None, ++ topk_length=topk_length, + ) + if is_thd: + grad_query = attn_bwd["dq"] +@@ -1534,15 +1614,51 @@ + total_comp = k_indexer.shape[0] + indexer_topk = indexer_topk_idxs.shape[-1] + ++ # The coeff=0 training path does not consume the teacher's segmented ++ # indexer LSE. Compact the complete attention set and reuse the exact ++ # same valid-prefix lengths in forward and backward. Preserve the ++ # legacy non-compact teacher path when loss_coeff > 0. ++ topk_length = None ++ if loss_coeff == 0: ++ topk_idxs, topk_length = _compact_flat_topk_idxs(topk_idxs) ++ + out_flat, lse, lse_indexer = _dsa_fwd_flash_mla( + query, + kv_full, + topk_idxs, + softmax_scale, + attn_sink=attn_sink, +- topk_length=None, +- indexer_topk=indexer_topk, ++ topk_length=topk_length, ++ indexer_topk=0 if topk_length is not None else indexer_topk, + ) ++ if topk_length is not None: ++ topk_idxs.clamp_min_(0) ++ if q_padding_mask is not None: ++ # Avoid cuDNN DSA's zero-tile backward path; dO/LSE for these ++ # rows are masked before the harmless placeholder is used. ++ topk_length.masked_fill_(q_padding_mask, 1) ++ ++ if loss_coeff == 0: ++ indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) ++ saved_grad_q_indexer = torch.zeros_like(q_indexer) ++ saved_grad_k_indexer = torch.zeros_like(k_indexer) ++ saved_grad_weights = torch.zeros_like(weights) ++ ctx.save_for_backward( ++ query, ++ kv_full, ++ attn_sink, ++ topk_idxs, ++ topk_length, ++ out_flat, ++ lse, ++ saved_grad_q_indexer, ++ saved_grad_k_indexer, ++ saved_grad_weights, ++ ) ++ ctx.softmax_scale = softmax_scale ++ ctx.has_topk_length = True ++ ctx.q_padding_mask = q_padding_mask ++ return out_flat.reshape(total_q, np_ * out_flat.shape[-1]), indexer_loss + + bwd_loss_coeff = loss_coeff * total_q / loss_divisor + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) +@@ -1674,11 +1790,13 @@ + saved_grad_k_indexer = torch.zeros_like(k_indexer) + saved_grad_weights = torch.zeros_like(weights) + ++ empty_topk_length = torch.empty(0, dtype=torch.int32, device=topk_idxs.device) + ctx.save_for_backward( + query, + kv_full, + attn_sink, + topk_idxs, ++ empty_topk_length, + out_flat, + lse, + saved_grad_q_indexer, +@@ -1686,6 +1804,8 @@ + saved_grad_weights, + ) + ctx.softmax_scale = softmax_scale ++ ctx.has_topk_length = False ++ ctx.q_padding_mask = q_padding_mask + + return out_flat.reshape(total_q, np_ * out_flat.shape[-1]), indexer_loss + +@@ -1698,14 +1818,19 @@ + kv_full, + attn_sink, + topk_idxs, ++ saved_topk_length, + out_flat, + lse, + saved_grad_q_indexer, + saved_grad_k_indexer, + saved_grad_weights, + ) = ctx.saved_tensors ++ topk_length = saved_topk_length if ctx.has_topk_length else None + + dO_flat = grad_output.reshape(query.shape[0], query.shape[1], out_flat.shape[-1]) ++ if topk_length is not None and ctx.q_padding_mask is not None: ++ dO_flat = dO_flat.masked_fill(ctx.q_padding_mask[:, None, None], 0) ++ lse = lse.masked_fill(ctx.q_padding_mask[:, None], 0) + attn_bwd = _DSA.sparse_attention_backward_wrapper( + query, + kv_full, +@@ -1715,7 +1840,7 @@ + attn_sink, + topk_idxs, + softmax_scale=ctx.softmax_scale, +- topk_length=None, ++ topk_length=topk_length, + ) + return ( + attn_bwd["dq"], +diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py +--- a/megatron/core/transformer/experimental_attention_variant/csa.py ++++ b/megatron/core/transformer/experimental_attention_variant/csa.py +@@ -2543,7 +2543,8 @@ + self.config, + ) + q_indexer_cp = rotate_activation(q_indexer_cp) +- weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) ++ with get_fp8_disabled_context(indexer.config): ++ weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) + weights_indexer_cp = weights_indexer_cp.squeeze(1) * (indexer.index_n_heads**-0.5) + + indexer_compressed_local, _ = indexer.compressor._forward_thd( +diff --git a/megatron/bridge/models/conversion/model_bridge.py b/megatron/bridge/models/conversion/model_bridge.py +--- a/megatron/bridge/models/conversion/model_bridge.py ++++ b/megatron/bridge/models/conversion/model_bridge.py +@@ -1080,7 +1080,7 @@ + _hf_import_cache: Dict[str, torch.Tensor] = {} + for task in self._with_progress_tracking(hf_to_megatron_tasks, description): + # None means megatron module not on current rank, skip if this task is not going to happen +- if task.megatron_module is None: ++ if task is None or task.megatron_module is None: + continue + # 1) Fetch source tensor(s) from HF state dict, with caching for grouped mappings + hf_param_key = str(task.mapping.hf_param) +@@ -1209,6 +1209,6 @@ + for task in conversion_tasks: + # None means megatron module not on current rank, skip if this task is not going to happen +- if task.megatron_module is None: ++ if task is None or task.megatron_module is None: + continue + hf_state_dict: Mapping[str, torch.Tensor] = hf_pretrained.state + if isinstance(task.mapping.hf_param, str): +@@ -1322,4 +1322,6 @@ + for task in self._with_progress_tracking(megatron_to_hf_tasks, "Converting to HuggingFace", show_progress): ++ if task is None: ++ continue + if isinstance(task.param_weight, DTensor): + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + uneven_dtensor_to_full_tensor, diff --git a/docker/patch/sglang/v0.5.15.post1.patch b/docker/patch/sglang/v0.5.15.post1.patch index f97f26003..1d66fb6e4 100644 --- a/docker/patch/sglang/v0.5.15.post1.patch +++ b/docker/patch/sglang/v0.5.15.post1.patch @@ -421,7 +421,7 @@ index bd9d7eafa1..f53a32abc0 100644 @app.post("/update_weight_version") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version( -@@ -1451,6 +1474,19 @@ async def load_lora_adapter_from_tensors( +@@ -1451,6 +1476,19 @@ async def load_lora_adapter_from_tensors( return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) diff --git a/docker/patch/sglang/v0.5.17.patch b/docker/patch/sglang/v0.5.17.patch new file mode 100644 index 000000000..2b8c09881 --- /dev/null +++ b/docker/patch/sglang/v0.5.17.patch @@ -0,0 +1,2454 @@ +diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +index 610267b..c592353 100644 +--- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py ++++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +@@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConf + from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, ++ TextConditioningOutput, + maybe_unpad_latents, + pad_text_embeddings_with_mask, + shard_rotary_emb_for_sp, +@@ -28,6 +29,9 @@ from sglang.multimodal_gen.runtime.utils.condition_expansion import ( + from sglang.multimodal_gen.runtime.utils.vision import resize + from sglang.multimodal_gen.utils import calculate_dimensions + ++_QWEN_PROMPT_TEMPLATE_START_IDX = 34 ++_QWEN_MAX_SEQUENCE_LENGTH = 512 ++ + + def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() +@@ -47,7 +51,10 @@ def qwen_image_preprocess_text(prompt): + + + def qwen_image_postprocess_text( +- outputs, _text_inputs, drop_idx=34, return_attention_mask=False ++ outputs, ++ _text_inputs, ++ drop_idx=_QWEN_PROMPT_TEMPLATE_START_IDX, ++ return_attention_mask=False, + ): + """Postprocess Qwen text embeddings. + +@@ -61,6 +68,25 @@ def qwen_image_postprocess_text( + ) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + conditioning = pad_text_embeddings_with_mask(split_hidden_states) ++ prompt_embeds = conditioning.prompt_embeds[:, :_QWEN_MAX_SEQUENCE_LENGTH] ++ prompt_embeds_mask = ( ++ conditioning.prompt_embeds_mask[:, :_QWEN_MAX_SEQUENCE_LENGTH] ++ if conditioning.prompt_embeds_mask is not None ++ else None ++ ) ++ prompt_seq_lens = ( ++ [ ++ min(int(seq_len), _QWEN_MAX_SEQUENCE_LENGTH) ++ for seq_len in conditioning.prompt_seq_lens ++ ] ++ if conditioning.prompt_seq_lens is not None ++ else None ++ ) ++ conditioning = TextConditioningOutput( ++ prompt_embeds, ++ prompt_embeds_mask, ++ prompt_seq_lens, ++ ) + if return_attention_mask: + return conditioning + return conditioning.prompt_embeds +@@ -179,6 +205,7 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig + dict( + padding=True, + truncation=True, ++ max_length=_QWEN_MAX_SEQUENCE_LENGTH + _QWEN_PROMPT_TEMPLATE_START_IDX, + ), + None, + ] +@@ -208,7 +235,10 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig + if tok_kwargs.get("max_length") is not None: + tok_kwargs["padding"] = "max_length" + else: +- tok_kwargs.setdefault("max_length", 1024) ++ tok_kwargs.setdefault( ++ "max_length", ++ _QWEN_MAX_SEQUENCE_LENGTH + _QWEN_PROMPT_TEMPLATE_START_IDX, ++ ) + tok_kwargs["padding"] = True + return tokenizer(prompts, **tok_kwargs) + +diff --git a/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py b/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py +index 9103832..15737cc 100644 +--- a/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py ++++ b/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py +@@ -11,7 +11,7 @@ from typing import Any, Callable + + from sglang.multimodal_gen.utils import StoreBoolean + +-_VALID_ROLLOUT_SDE_TYPES = ("sde", "cps", "ode") ++_VALID_ROLLOUT_SDE_TYPES = ("sde", "cps", "ode", "dance") + + + @dataclass +diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py +index c5636c5..7ddbad7 100644 +--- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py ++++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py +@@ -241,6 +241,16 @@ class SamplingParams: + # 0-indexed denoising-loop step filters; None = all steps. + rollout_sde_step_indices: list[int] | None = None + rollout_return_step_indices: list[int] | None = None ++ # Driver-supplied rollout reproducibility data. These fields are consumed by ++ # the diffusion post-training path; prepare_request copies them into Req so ++ # the scheduler, denoising and latent-preparation stages see the same recipe ++ # the trainer will replay. ++ initial_noise_group_ids: list[str] | None = None ++ initial_noise_latent_shape: list[int] | None = None ++ initial_noise_seed: int | None = None ++ denoise_seeds: list[str] | None = None ++ sigmas: list[float] | None = None ++ timesteps: list[float] | None = None + # if True, disallow user params to override subclass-defined protected fields + no_override_protected_fields: bool = field( + default=False, metadata={"batch_sig_exclude": True} +diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +index 45b3676..236e19b 100644 +--- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py ++++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +@@ -37,6 +37,17 @@ class UpdateWeightFromTensorCheckerReqInput: + expected_named_tensors_sha256: dict[str, str] + + ++@dataclass ++class SetLoraFromTensorReqInput: ++ """Request to install a LoRA adapter from an in-memory tensor payload.""" ++ ++ serialized_tensors: str | bytes ++ lora_name: str = "default" ++ target: str = "transformer" ++ strength: float = 1.0 ++ merge_mode: str | None = None ++ ++ + @dataclass + class GetWeightsChecksumReqInput: + """Compute SHA-256 checksum of loaded module weights for verification.""" +diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py +index 554feea..251dde7 100644 +--- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py ++++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py +@@ -6,6 +6,7 @@ from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + GetWeightsChecksumReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, ++ SetLoraFromTensorReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, +@@ -138,6 +139,51 @@ async def update_weights_from_tensor_checker(request: Request): + ) + + ++@router.post("/set_lora_from_tensor") ++async def set_lora_from_tensor(request: Request): ++ """Install a LoRA adapter from base64(torch.save(dict[str, Tensor])).""" ++ import base64 ++ ++ body = await request.json() ++ serialized = body.get("serialized_tensors") ++ if serialized is None: ++ return orjson_response( ++ {"success": False, "message": "serialized_tensors is required"}, ++ status_code=400, ++ ) ++ ++ req = SetLoraFromTensorReqInput( ++ serialized_tensors=( ++ base64.b64decode(serialized) if isinstance(serialized, str) else serialized ++ ), ++ lora_name=body.get("lora_name", "default"), ++ target=body.get("target", "transformer"), ++ strength=float(body.get("strength", 1.0)), ++ merge_mode=body.get("merge_mode"), ++ ) ++ ++ try: ++ response = await async_scheduler_client.forward(req) ++ except Exception as e: ++ return orjson_response({"success": False, "message": str(e)}, status_code=500) ++ ++ if response.error is not None: ++ return orjson_response( ++ {"success": False, "message": str(response.error)}, ++ status_code=400, ++ ) ++ result = response.output or {} ++ success = bool(result.get("success", False)) ++ return orjson_response( ++ { ++ "success": success, ++ "message": result.get("message", "Unknown status"), ++ "adapted_layers": result.get("adapted_layers"), ++ }, ++ status_code=200 if success else 400, ++ ) ++ ++ + @router.post("/get_weights_checksum") + async def get_weights_checksum(request: Request): + """Return SHA-256 checksum of each requested module's weights.""" +diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +index 1494f58..d3cb820 100644 +--- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py ++++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +@@ -759,6 +759,19 @@ def prepare_request( + VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, + ) + sampling_params.apply_request_extra(req) ++ for name in ( ++ "initial_noise_group_ids", ++ "initial_noise_latent_shape", ++ "initial_noise_seed", ++ "denoise_seeds", ++ "sigmas", ++ "timesteps", ++ ): ++ value = getattr(sampling_params, name, None) ++ if value is not None: ++ if name == "timesteps": ++ value = torch.as_tensor(value, dtype=torch.float32) ++ setattr(req, name, value) + if getattr(sampling_params, "max_sequence_length", None) is not None: + req.max_sequence_length = sampling_params.max_sequence_length + +diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py +index cf64593..d0efdf3 100644 +--- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py ++++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py +@@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + GetWeightsChecksumReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, ++ SetLoraFromTensorReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, +@@ -135,6 +136,7 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag + ShutdownReq: self._handle_shutdown, + ReleaseRealtimeSessionReq: self._handle_release_realtime_session, + GetDisaggStatsReq: self._handle_get_disagg_stats, ++ SetLoraFromTensorReqInput: self._handle_set_lora_from_tensor, + UpdateWeightFromDiskReqInput: self._handle_update_weights_from_disk, + UpdateWeightFromTensorReqInput: self._handle_update_weights_from_tensor, + UpdateWeightFromTensorCheckerReqInput: ( +diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +index 1f25398..2ecef7d 100644 +--- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py ++++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +@@ -350,21 +350,23 @@ class FlowMatchEulerDiscreteScheduler( + sigmas_array = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas_array) + +- # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of +- # "exponential" or "linear" type is applied +- if self.config.use_dynamic_shifting: +- assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" +- sigmas_array = self.time_shift(mu, 1.0, sigmas_array) +- else: +- sigmas_array = ( +- self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array) +- ) ++ # 2. Perform timestep shifting. Driver-supplied sigmas are already in the ++ # exact sigma space the trainer will replay, so do not shift or ++ # stretch them again. ++ if sigmas is None: ++ if self.config.use_dynamic_shifting: ++ assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" ++ sigmas_array = self.time_shift(mu, 1.0, sigmas_array) ++ else: ++ sigmas_array = ( ++ self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array) ++ ) + +- # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value +- if self.config.shift_terminal: +- sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) +- sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) +- sigmas_array = sigmas_tensor.numpy() ++ # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value ++ if self.config.shift_terminal: ++ sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) ++ sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) ++ sigmas_array = sigmas_tensor.numpy() + + # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules + if self.config.use_karras_sigmas: +diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +index 8f3b26a..fbceed3 100644 +--- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py ++++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +@@ -722,6 +722,21 @@ class LoRAPipeline(ComposedPipelineBase): + if adapter_config.get("lora_alpha") is not None: + adapter_lora_alpha = int(adapter_config["lora_alpha"]) + ++ self._ingest_lora_state_dict( ++ lora_state_dict, ++ lora_nickname, ++ adapter_lora_alpha=adapter_lora_alpha, ++ ) ++ self.loaded_adapter_paths[lora_nickname] = lora_path ++ logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) ++ ++ def _ingest_lora_state_dict( ++ self, ++ lora_state_dict: dict[str, torch.Tensor], ++ lora_nickname: str, ++ adapter_lora_alpha: int | None = None, ++ ) -> None: ++ """Map adapter keys to SGLang DiT names and stage them on device.""" + if lora_nickname in self.lora_adapters: + self.lora_adapters[lora_nickname].clear() + +@@ -764,9 +779,91 @@ class LoRAPipeline(ComposedPipelineBase): + f"Dit target weight name {target_name} already exists in lora_adapters[{lora_nickname}]" + ) + self.lora_adapters[lora_nickname][target_name] = weight.to(self.device) +- self.loaded_adapter_paths[lora_nickname] = lora_path + self.loaded_adapter_alphas[lora_nickname] = adapter_lora_alpha +- logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) ++ ++ def set_lora_from_tensors( ++ self, ++ lora_nickname: str, ++ named_tensors: dict[str, torch.Tensor], ++ target: str = "transformer", ++ strength: float = 1.0, ++ merge_mode: str | None = None, ++ ) -> int: ++ """Apply a freshly supplied in-memory LoRA adapter.""" ++ if target not in self.VALID_TARGETS: ++ raise ValueError( ++ f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}" ++ ) ++ ++ if not self.lora_initialized: ++ with self._temporarily_disable_offload( ++ target="all", use_module_names_only=True ++ ): ++ self.convert_to_lora_layers() ++ ++ lora_state_dict = normalize_lora_state_dict( ++ dict(named_tensors), logger=logger ++ ) ++ self._ingest_lora_state_dict(lora_state_dict, lora_nickname) ++ tensor_path = f"" ++ self.loaded_adapter_paths[lora_nickname] = tensor_path ++ ++ target_modules, error = self._get_target_lora_layers(target) ++ if error: ++ raise ValueError(f"set_lora_from_tensors: {error}") ++ ++ merge_mode = self._resolve_lora_merge_mode(None, merge_mode) ++ merge_weights_by_module = { ++ module_name: self._should_merge_lora_for_layers( ++ module_name, lora_layers_dict, merge_mode ++ ) ++ for module_name, lora_layers_dict in target_modules ++ } ++ ++ if self._needs_lora_weight_update_context( ++ target_modules, merge_weights_by_module ++ ): ++ weight_update_context = self._temporarily_disable_offload( ++ target_modules=target_modules ++ ) ++ else: ++ weight_update_context = nullcontext() ++ ++ adapted_count = 0 ++ rank = dist.get_rank() ++ with weight_update_context: ++ for module_name, lora_layers_dict in target_modules: ++ effective_merge_weights = merge_weights_by_module[module_name] ++ count = self._apply_lora_to_layers( ++ lora_layers_dict, ++ [lora_nickname], ++ [tensor_path], ++ rank, ++ [strength], ++ clear_existing=True, ++ merge_weights=effective_merge_weights, ++ ) ++ adapted_count += count ++ self.cur_adapter_name[module_name] = lora_nickname ++ self.cur_adapter_path[module_name] = tensor_path ++ self.is_lora_merged[module_name] = effective_merge_weights ++ self.cur_adapter_strength[module_name] = strength ++ self.cur_adapter_config[module_name] = ( ++ [lora_nickname], ++ [strength], ++ ) ++ ++ logger.info( ++ "Rank %d: in-memory LoRA adapter %s applied to %d layers " ++ "(target=%s, strength=%.2f, merge_mode=%s)", ++ rank, ++ lora_nickname, ++ adapted_count, ++ target, ++ strength, ++ merge_mode, ++ ) ++ return adapted_count + + def set_lora( + self, +diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +index bc1cb59..4a74aa7 100644 +--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py ++++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +@@ -6,6 +6,7 @@ Denoising stage for diffusion pipelines. + """ + + import gc ++import hashlib + import inspect + import math + import time +@@ -135,6 +136,30 @@ from sglang.multimodal_gen.runtime.utils.torch_compile import ( + from sglang.multimodal_gen.utils import dict_to_3d_list + + logger = init_logger(__name__) ++_MAX_TORCH_SEED = (1 << 63) - 1 ++ ++ ++def _resolve_base_seed(batch) -> int | None: ++ seed = getattr(batch, "seed", None) ++ if seed is None: ++ seed = getattr(getattr(batch, "sampling_params", None), "seed", None) ++ return int(seed) if seed is not None else None ++ ++ ++def _make_step_generators( ++ base_seed: int, ++ step_index: int, ++ denoise_seeds: list[str], ++) -> list[torch.Generator]: ++ generators = [] ++ for seed_key in denoise_seeds: ++ payload = f"{int(base_seed)}::step::{int(step_index)}::sample::{str(seed_key)}".encode("utf-8") ++ digest = hashlib.blake2b(payload, digest_size=8).digest() ++ seed = int.from_bytes(digest, byteorder="big", signed=False) % _MAX_TORCH_SEED ++ generator = torch.Generator(device="cpu") ++ generator.manual_seed(seed) ++ generators.append(generator) ++ return generators + + + def _ensure_tensor_model_output(model_output): +@@ -1229,6 +1254,15 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): + below; mirror them in the override if those markers are needed. + """ + use_nvtx = self.current_use_nvtx ++ denoise_seeds = getattr(batch, "denoise_seeds", None) ++ base_seed = _resolve_base_seed(batch) ++ if getattr(batch, "rollout", False) and denoise_seeds is not None and base_seed is not None: ++ ctx.extra_step_kwargs["generator"] = _make_step_generators( ++ base_seed, ++ int(step.step_index), ++ [str(seed) for seed in denoise_seeds], ++ ) ++ + # 1. Prepare latent inputs in the model's compute dtype. + latent_model_input = ctx.latents.to(ctx.target_dtype) + if batch.image_latent is not None: +diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +index f9a28a2..40c705f 100644 +--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py ++++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +@@ -5,6 +5,7 @@ + Latent preparation stage for diffusion pipelines. + """ + ++import hashlib + from dataclasses import dataclass + from typing import Any + +@@ -26,6 +27,48 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs + from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + + logger = init_logger(__name__) ++_MAX_TORCH_SEED = (1 << 63) - 1 ++ ++ ++def _derive_group_seed(base_seed: int, group_id: str) -> int: ++ payload = f"{int(base_seed)}::{str(group_id)}".encode("utf-8") ++ digest = hashlib.blake2b(payload, digest_size=8).digest() ++ return int.from_bytes(digest, byteorder="big", signed=False) % (_MAX_TORCH_SEED + 1) ++ ++ ++def _regen_initial_noise( ++ noise_group_ids: list[str], ++ base_seed: int, ++ latent_shape: list[int] | tuple[int, ...], ++ device, ++ dtype, ++): ++ tensors = [] ++ cache = {} ++ for raw_group_id in noise_group_ids: ++ group_id = str(raw_group_id) ++ noise = cache.get(group_id) ++ if noise is None: ++ generator = torch.Generator(device=torch.device("cpu")) ++ generator.manual_seed(_derive_group_seed(base_seed, group_id)) ++ noise = torch.randn( ++ tuple(int(v) for v in latent_shape), ++ generator=generator, ++ device="cpu", ++ dtype=torch.float32, ++ ) ++ cache[group_id] = noise ++ tensors.append(noise) ++ return torch.stack(tensors, dim=0).to(device=device, dtype=dtype) ++ ++ ++def _driver_xt_recipe(batch): ++ group_ids = getattr(batch, "initial_noise_group_ids", None) ++ latent_shape = getattr(batch, "initial_noise_latent_shape", None) ++ base_seed = getattr(batch, "initial_noise_seed", None) ++ if group_ids and latent_shape is not None and base_seed is not None: ++ return [str(v) for v in group_ids], [int(v) for v in latent_shape], int(base_seed) ++ return None + + + @dataclass(frozen=True) +@@ -127,6 +170,7 @@ class LatentPreparationStage(PipelineStage): + latents = batch.latents + height = batch.height + width = batch.width ++ driver_xt_recipe = _driver_xt_recipe(batch) + + # TODO(will): remove this once we add input/output validation for stages + if self.requires_batch_height_width(batch, server_args) and ( +@@ -142,7 +186,34 @@ class LatentPreparationStage(PipelineStage): + ) + + # Generate or use provided latents +- if latents is None: ++ if driver_xt_recipe is not None: ++ group_ids, latent_shape, base_seed = driver_xt_recipe ++ spec = self.get_latent_preparation_spec( ++ batch, server_args, len(group_ids), latent_num_frames, device ++ ) ++ raw_latents = _regen_initial_noise( ++ group_ids, ++ base_seed, ++ latent_shape, ++ device=spec.device, ++ dtype=spec.dtype, ++ ) ++ ++ latent_ids = ( ++ server_args.pipeline_config.maybe_prepare_latent_ids(raw_latents) ++ if spec.prepare_latent_ids ++ else None ++ ) ++ ++ if latent_ids is not None: ++ batch.latent_ids = latent_ids.to(device=device) ++ ++ latents = raw_latents ++ if spec.pack_latents: ++ latents = server_args.pipeline_config.maybe_pack_latents( ++ raw_latents, len(group_ids), batch ++ ) ++ elif latents is None: + spec = self.get_latent_preparation_spec( + batch, server_args, batch_size, latent_num_frames, device + ) +@@ -202,7 +273,8 @@ class LatentPreparationStage(PipelineStage): + indexed_batches = group + group_batches = [batch for _, batch in indexed_batches] + if len(group_batches) == 1 or any( +- batch.latents is not None for batch in group_batches ++ batch.latents is not None or _driver_xt_recipe(batch) is not None ++ for batch in group_batches + ): + for index, batch in indexed_batches: + results[index] = self(batch, server_args) +diff --git a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py +index f652059..d2e6eb7 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py +@@ -7,6 +7,7 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_ch + from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + iter_materialized_weights, + ) ++from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch + from sglang.multimodal_gen.runtime.post_training.tensor_update_checker import ( + TensorUpdateChecker, + ) +@@ -19,12 +20,58 @@ from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions + + if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( ++ SetLoraFromTensorReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, + ) + + + class GPUWorkerPostTrainingMixin: ++ def set_lora_from_tensors( ++ self, ++ req: SetLoraFromTensorReqInput, ++ ) -> OutputBatch: ++ """Install a LoRA adapter from a torch.save blob of tensor weights.""" ++ import io ++ ++ import torch ++ ++ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline ++ ++ if not isinstance(self.pipeline, LoRAPipeline): ++ return OutputBatch(error="Lora is not enabled") ++ ++ try: ++ named_tensors = torch.load( ++ io.BytesIO(req.serialized_tensors), ++ map_location="cpu", ++ weights_only=True, ++ ) ++ adapted = self.pipeline.set_lora_from_tensors( ++ req.lora_name, ++ named_tensors, ++ target=req.target, ++ strength=req.strength, ++ merge_mode=req.merge_mode, ++ ) ++ if adapted <= 0: ++ return OutputBatch( ++ error=f"LoRA adapter {req.lora_name} did not match any layer" ++ ) ++ except Exception as e: ++ return OutputBatch(error=str(e)) ++ ++ return OutputBatch( ++ output={ ++ "success": True, ++ "message": ( ++ f"Applied in-memory LoRA {req.lora_name} to " ++ f"{adapted} layers of {req.target}" ++ ), ++ "adapted_layers": adapted, ++ } ++ ) ++ + def update_weights_from_disk( + self, + model_path: str, +@@ -135,7 +182,7 @@ class GPUWorkerPostTrainingMixin: + + def get_weights_checksum( + self, module_names: list[str] | None = None +- ) -> dict[str, str]: ++ ) -> dict[str, str | int]: + if not self.pipeline: + return {"error": "Pipeline is not initialized"} + +@@ -143,15 +190,16 @@ class GPUWorkerPostTrainingMixin: + names = module_names if module_names is not None else list(all_modules.keys()) + + checksums: dict[str, str] = {} ++ tensor_count = 0 + for name in names: + module = all_modules.get(name) + if module is None: + checksums[name] = "not_found" + continue +- checksums[name] = compute_weights_checksum( +- iter_materialized_weights(module) +- ) +- return checksums ++ materialized_weights = list(iter_materialized_weights(module)) ++ tensor_count += len(materialized_weights) ++ checksums[name] = compute_weights_checksum(materialized_weights) ++ return {**checksums, "tensor_count": tensor_count} + + def _select_rank_scoped_payload( + self, +diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py +index 025643b..8dfe997 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py +@@ -6,6 +6,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa + + + class SchedulerPostTrainingMixin: ++ def _handle_set_lora_from_tensor(self, reqs: List[Any]) -> OutputBatch: ++ req = reqs[0] ++ return self.worker.set_lora_from_tensors(req) ++ + def _handle_update_weights_from_disk(self, reqs: List[Any]) -> OutputBatch: + req = reqs[0] + success, message = self.worker.update_weights_from_disk( +diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +index 3777ad4..c4e91ea 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +@@ -1,7 +1,9 @@ + # SPDX-License-Identifier: Apache-2.0 + """Flow-matching rollout step utilities for log-prob computation.""" + ++import hashlib + import math ++import os + from typing import Any, Union + + import torch +@@ -18,6 +20,25 @@ from sglang.multimodal_gen.runtime.post_training.scheduler_rl_debug_mixin import + ) + + _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) ++_MAX_TORCH_SEED = (1 << 63) - 1 ++ ++ ++def _resolve_base_seed(batch) -> int | None: ++ seed = getattr(batch, "seed", None) ++ if seed is None: ++ seed = getattr(getattr(batch, "sampling_params", None), "seed", None) ++ return int(seed) if seed is not None else None ++ ++ ++def _resolve_fallback_seed(batch) -> int: ++ base_seed = _resolve_base_seed(batch) ++ denoise_seeds = getattr(batch, "denoise_seeds", None) ++ sample_key = str(denoise_seeds[0]) if denoise_seeds else None ++ if base_seed is not None and sample_key is not None: ++ payload = f"{int(base_seed)}::fallback::sample::{sample_key}".encode("utf-8") ++ digest = hashlib.blake2b(payload, digest_size=8).digest() ++ return int.from_bytes(digest, byteorder="big", signed=False) % _MAX_TORCH_SEED ++ return int.from_bytes(os.urandom(8), byteorder="big") % _MAX_TORCH_SEED + + + class SchedulerRLMixin(SchedulerRLDebugMixin): +@@ -86,7 +107,15 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + B = local_shape[0] + if isinstance(generator, torch.Generator): + assert B == 1, "Generator must be a list if batch size is not 1" +- generator = [generator] ++ per_request_generator = getattr(batch, "_relax_noise_gen", None) ++ if per_request_generator is None: ++ per_request_generator = torch.Generator(device=device) ++ per_request_generator.manual_seed(_resolve_fallback_seed(batch)) ++ try: ++ batch._relax_noise_gen = per_request_generator ++ except AttributeError: ++ pass ++ generator = [per_request_generator] + else: + assert ( + len(generator) == B +@@ -96,11 +125,14 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + rollout_session_data, rollout_session_data.latents_shape, device, dtype + ) + for i in range(B): +- torch.randn( +- rollout_session_data.latents_shape, +- out=buffer[i : i + 1], +- generator=generator[i], +- ) ++ gen = generator[i] ++ gen_device = getattr(gen, "device", None) ++ sample_shape = tuple(buffer[i].shape) ++ if gen is not None and gen_device is not None and gen_device.type != buffer.device.type: ++ tmp = torch.randn(sample_shape, generator=gen, dtype=dtype, device=gen_device) ++ buffer[i].copy_(tmp) ++ else: ++ torch.randn(sample_shape, out=buffer[i], generator=gen) + + sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp( + batch=batch, latents=buffer +@@ -128,6 +160,7 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + 1. ``"sde"``: Standard stochastic differential equation transition (Gaussian). + 2. ``"cps"``: Coupled Particle Sampling. + 3. ``"ode"``: Deterministic ODE step (no diffusion noise). ++ 4. ``"dance"``: FlowGRPO dance-style transition used by Relax diffusion RL. + """ + rollout_session_data = self._get_rollout_session_data(batch) + sde_type = batch.rollout_sde_type +@@ -234,6 +267,34 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + log_prob_no_const + ), "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" + ++ elif effective_sde_type == "dance": ++ model_output = model_output.float() ++ sample = sample.float() ++ variance_noise = self._rollout_variance_noise( ++ batch, model_output, generator ++ ) ++ full_variance_noise = rollout_session_data.noise_buffer ++ std_dev_t = current_sigma.new_tensor(noise_level) ++ noise_std_dev = std_dev_t * torch.sqrt(-1 * dt) ++ prev_sample_mean = ( ++ sample * (1 + std_dev_t**2 / (2 * current_sigma) * dt) ++ + model_output ++ * (1 + std_dev_t**2 * (1 - current_sigma) / (2 * current_sigma)) ++ * dt ++ ) ++ prev_sample = prev_sample_mean + variance_noise * noise_std_dev ++ if tuple(variance_noise.shape) != tuple(full_variance_noise.shape): ++ raise RuntimeError( ++ "rollout_sde_type='dance' does not support sequence-parallel rollout: " ++ f"the per-step log-prob is computed over the unsharded noise " ++ f"{tuple(full_variance_noise.shape)} while this rank only owns " ++ f"{tuple(variance_noise.shape)}, so the cross-rank reduction would " ++ "over-count it by the SP degree. Run the diffusion engine with " ++ "--rollout-num-gpus-per-engine 1 until the reduction convention " ++ "is confirmed against the live engine." ++ ) ++ log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2) ++ + else: + raise ValueError(f"Unsupported sde_type: {sde_type}") + +diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +index 7c470d0..2b50e0e 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +@@ -493,6 +493,13 @@ class WeightsUpdater: + lora_rank=lora_rank, + ) + ++ if getattr(self.pipeline, "lora_initialized", False): ++ return False, ( ++ "Refusing update_weights_from_tensor: this pipeline has LoRA " ++ "layers installed. Use adapter-only sync (/set_lora_from_tensor) " ++ "for a LoRA run." ++ ) ++ + if target_modules is None: + target_modules = [_DEFAULT_TENSOR_TARGET_MODULE] + try: +diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py +index 240e20de..eda3406c 100644 +--- a/python/sglang/srt/disaggregation/decode.py ++++ b/python/sglang/srt/disaggregation/decode.py +@@ -21,6 +21,7 @@ Life cycle of a request in the decode server + from __future__ import annotations + + import logging ++import os + import time + from collections import deque + from dataclasses import dataclass +@@ -680,6 +681,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + def release_memory_occupation(self): + self.queue.clear() + self.retracted_queue.clear() ++ self.pending_reqs.clear() + if hasattr(self.kv_manager, "deregister_buffer_to_engine"): + self.kv_manager.deregister_buffer_to_engine() + +@@ -769,6 +771,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + ) + ++ bootstrap_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + for decode_req, poll in zip(self.queue, polls): + if poll is None: + continue +@@ -776,7 +783,21 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + continue + + if poll == KVPoll.Bootstrapping: +- pass ++ entry_time = decode_req.req.time_stats.decode_prealloc_queue_entry_time ++ if entry_time > 0 and now - entry_time > bootstrap_timeout: ++ error_message = ( ++ f"Decode prealloc timeout for request rank={self.tp_rank} " ++ f"{decode_req.req.rid=} {decode_req.req.bootstrap_room=} " ++ f"after {bootstrap_timeout}s" ++ ) ++ logger.error(error_message) ++ prepare_abort( ++ decode_req.req, ++ error_message, ++ status_code=HTTPStatus.GATEWAY_TIMEOUT, ++ ) ++ if self.scheduler.metrics_reporter.enable_metrics: ++ self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() + elif poll == KVPoll.WaitingForInput: + decode_req.waiting_for_input = True + decode_req.req.time_stats.set_bootstrap_done_time() +@@ -1995,6 +2016,11 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + else: + polls = self._poll_with_metadata_gate() + ++ transfer_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + transferred_reqs = [] + indices_to_remove = set() + for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): +@@ -2071,7 +2097,17 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + KVPoll.WaitingForInput, + KVPoll.Transferring, + ]: +- pass ++ entry_time = decode_req.req.time_stats.decode_transfer_queue_entry_time ++ if entry_time > 0 and now - entry_time > transfer_timeout: ++ logger.error( ++ "Decode transfer timeout for request rank=%s rid=%s room=%s " ++ "after %ss", ++ self.tp_rank, ++ decode_req.req.rid, ++ decode_req.req.bootstrap_room, ++ transfer_timeout, ++ ) ++ decode_req.kv_receiver.abort() + else: + raise ValueError(f"Unexpected poll case: {poll}") + +@@ -2097,6 +2133,9 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + + def release_memory_occupation(self): + """Clean up in-flight transfers before releasing GPU memory.""" ++ for decode_req in self.queue: ++ if decode_req.kv_receiver is not None: ++ decode_req.kv_receiver.abort() + self.queue.clear() + + def resume_memory_occupation(self): +@@ -2320,6 +2359,11 @@ class SchedulerDisaggregationDecodeMixin: + resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() + self.waiting_queue.extend(resumed_reqs) + if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: ++ transferred_reqs = self.disagg_decode_transfer_queue.pop_transferred() ++ if self.enable_hisparse: ++ for req in transferred_reqs: ++ self.hisparse_coordinator.admit_request_direct(req) ++ self.waiting_queue.extend(transferred_reqs) + # if there are still retracted requests, we do not allocate new requests + return + +diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py +index 1907e0ee..14048ad0 100644 +--- a/python/sglang/srt/disaggregation/mooncake/conn.py ++++ b/python/sglang/srt/disaggregation/mooncake/conn.py +@@ -1050,7 +1050,7 @@ class MooncakeKVManager(CommonKVManager): + for i, dst_aux_ptr in enumerate(dst_aux_ptrs): + length = prefill_aux_item_lens[i] + src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index +- dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index ++ dst_addr = dst_aux_ptr + length * req.dst_aux_index + transfer_blocks.append((src_addr, dst_addr, length)) + + return self._transfer_data(req.mooncake_session_id, transfer_blocks) +@@ -1714,12 +1714,6 @@ class MooncakeKVManager(CommonKVManager): + if ret != 0: + with self.session_lock: + self.session_failures[req.mooncake_session_id] += 1 +- # Failures should never happen if the session is not dead, if the session fails once, mark it as failed +- if self.session_failures[req.mooncake_session_id] >= 1: +- self.failed_sessions.add(req.mooncake_session_id) +- logger.error( +- f"Session {req.mooncake_session_id} failed." +- ) + self.record_failure( + kv_chunk.room, + f"Failed to send kv chunk of {kv_chunk.room} to " +diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py +index 5642c64e..599b9815 100644 +--- a/python/sglang/srt/disaggregation/prefill.py ++++ b/python/sglang/srt/disaggregation/prefill.py +@@ -21,6 +21,8 @@ from __future__ import annotations + + import hashlib + import logging ++import os ++import time + from array import array + from collections import deque + from http import HTTPStatus +@@ -416,6 +418,11 @@ class PrefillBootstrapQueue: + self.scheduler.attn_tp_cpu_group, + ) + ++ bootstrap_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + for i, (req, poll) in enumerate(zip(self.queue, polls)): + if poll is None: + continue +@@ -425,6 +432,27 @@ class PrefillBootstrapQueue: + indices_to_remove.add(i) + failed_reqs.append(req) + elif poll == KVPoll.Bootstrapping: ++ entry_time = req.time_stats.prefill_bootstrap_queue_entry_time ++ if entry_time > 0 and now - entry_time > bootstrap_timeout: ++ error_message = ( ++ f"Prefill bootstrap timed out after {now - entry_time:.1f}s " ++ f"for request rank={self.tp_rank} " ++ f"{req.rid=} {req.bootstrap_room=}" ++ ) ++ logger.error(error_message) ++ prepare_abort( ++ req, ++ error_message, ++ status_code=HTTPStatus.GATEWAY_TIMEOUT, ++ ) ++ self.scheduler.output_streamer.stream_output( ++ [req], req.return_logprob ++ ) ++ indices_to_remove.add(i) ++ failed_reqs.append(req) ++ if self.scheduler.metrics_reporter.enable_metrics: ++ self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() ++ continue + if ( + req.prefill_attempt_count + < self.scheduler.server_args.optimistic_prefill_attempts +@@ -832,6 +860,11 @@ class SchedulerDisaggregationPrefillMixin: + self.attn_tp_cpu_group, + ) + ++ transfer_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + undone_reqs: List[Req] = [] + # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue + for req, poll in zip(self.disagg_prefill_inflight_queue, polls): +@@ -866,7 +899,26 @@ class SchedulerDisaggregationPrefillMixin: + + if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + # todo: set Transferring correctly in backend +- undone_reqs.append(req) ++ entry_time = req.time_stats.prefill_transfer_queue_entry_time ++ if entry_time > 0 and now - entry_time > transfer_timeout: ++ error_message = ( ++ f"Prefill transfer timed out after {now - entry_time:.1f}s " ++ f"(state={poll}) for request rank={self.ps.tp_rank} " ++ f"{req.rid=} {req.bootstrap_room=}" ++ ) ++ logger.error(error_message) ++ release_kv_cache(req, self.tree_cache) ++ prepare_abort( ++ req, ++ error_message, ++ status_code=HTTPStatus.GATEWAY_TIMEOUT, ++ ) ++ req.disagg_kv_sender.clear() ++ done_reqs.append(req) ++ if self.metrics_reporter.enable_metrics: ++ self.metrics_collector.increment_transfer_failed_reqs() ++ else: ++ undone_reqs.append(req) + elif poll == KVPoll.Success: # transfer done + if not isinstance(req.finished_reason, FINISH_ABORT): + req.finished_reason = FINISH_LENGTH(length=0) +diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py +index 6eb0ad1a..d27815f3 100644 +--- a/python/sglang/srt/entrypoints/engine.py ++++ b/python/sglang/srt/entrypoints/engine.py +@@ -73,6 +73,7 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqInput, + MultimodalDataInputFormat, + OpenSessionReqInput, ++ PostProcessWeightsReqInput, + ProfileReq, + ProfileReqType, + ReleaseMemoryOccupationReqInput, +@@ -81,6 +82,7 @@ from sglang.srt.managers.io_struct import ( + RpcReqInput, + RpcReqOutput, + UnloadLoRAAdapterReqInput, ++ UpdateLoRAFromDistributedReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -1429,6 +1431,20 @@ class Engine(EngineScoreMixin, EngineBase): + self.tokenizer_manager.update_weights_from_ipc(obj, None) + ) + ++ def post_process_weights( ++ self, ++ restore_weights_before_load: bool = False, ++ post_process_quantization: bool = False, ++ ): ++ """Optional post-processing for updated weights, e.g. quantization packing.""" ++ obj = PostProcessWeightsReqInput( ++ restore_weights_before_load=restore_weights_before_load, ++ post_process_quantization=post_process_quantization, ++ ) ++ return self.loop.run_until_complete( ++ self.tokenizer_manager.post_process_weights(obj, None) ++ ) ++ + def get_weights_by_name(self, name: str, truncate_size: int = 100): + """Get weights by parameter name.""" + obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) +@@ -1474,6 +1490,36 @@ class Engine(EngineScoreMixin, EngineBase): + self.tokenizer_manager.load_lora_adapter_from_tensors(lora_req, None) + ) + ++ def update_lora_from_distributed( ++ self, ++ lora_name: str, ++ names, ++ dtypes, ++ shapes, ++ config_dict: Dict, ++ group_name: str, ++ pinned: bool = False, ++ ): ++ """Load/refresh a LoRA adapter whose tensors arrive over an NCCL group. ++ ++ In-memory counterpart to ``load_lora_adapter`` (disk) and ++ ``load_lora_adapter_from_tensors`` (serialized blob): the tensors are ++ broadcast ``src=0`` on ``group_name`` (the weight-update group), so ++ only metadata travels through this HTTP call. ++ """ ++ obj = UpdateLoRAFromDistributedReqInput( ++ lora_name=lora_name, ++ config_dict=config_dict, ++ names=names, ++ dtypes=dtypes, ++ shapes=shapes, ++ group_name=group_name, ++ pinned=pinned, ++ ) ++ return self.loop.run_until_complete( ++ self.tokenizer_manager.update_lora_from_distributed(obj, None) ++ ) ++ + def load_lora_adapter(self, lora_name: str, lora_path: str, pinned: bool = False): + """Load a new LoRA adapter without re-launching the engine.""" + +diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py +index ac38b466..ea54bac1 100644 +--- a/python/sglang/srt/entrypoints/http_server.py ++++ b/python/sglang/srt/entrypoints/http_server.py +@@ -132,6 +132,7 @@ from sglang.srt.managers.io_struct import ( + OpenSessionReqInput, + ParseFunctionCallReq, + PauseGenerationReqInput, ++ PostProcessWeightsReqInput, + ProfileReq, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, +@@ -140,6 +141,7 @@ from sglang.srt.managers.io_struct import ( + SetInternalStateReq, + SlowDownReqInput, + UnloadLoRAAdapterReqInput, ++ UpdateLoRAFromDistributedReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -762,10 +764,8 @@ async def model_info(): + @app.get("/weight_version") + async def weight_version(): + """Get the current weight version.""" +- raise HTTPException( +- status_code=404, +- detail="Endpoint '/get_weight_version' or '/weight_version' is deprecated. Please use '/model_info' instead.", +- ) ++ result = await model_info() ++ return {"weight_version": result.get("weight_version", None)} + + + @app.get("/get_server_info") +@@ -782,9 +782,18 @@ async def get_server_info(): + async def server_info(): + """Get the server information.""" + # Returns internal states per DP. +- internal_states: List[Dict[Any, Any]] = ( +- await _global_state.tokenizer_manager.get_internal_state() +- ) ++ server_info_timeout = float(os.environ.get("SGLANG_SERVER_INFO_TIMEOUT", "2")) ++ try: ++ internal_states: List[Dict[Any, Any]] = await asyncio.wait_for( ++ _global_state.tokenizer_manager.get_internal_state(), ++ timeout=server_info_timeout, ++ ) ++ except asyncio.TimeoutError: ++ logger.warning( ++ "Timed out getting internal state for /server_info after %.1fs; returning empty internal_states", ++ server_info_timeout, ++ ) ++ internal_states = [] + + server_args = _global_state.tokenizer_manager.server_args + +@@ -1409,6 +1418,22 @@ async def update_weights_from_ipc( + return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) + + ++@app.post("/post_process_weights") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def post_process_weights( ++ obj: Annotated[PostProcessWeightsReqInput, Body()], request: Request ++): ++ """Optional post-processing for updated weights, e.g. quantization packing.""" ++ success, message = await _global_state.tokenizer_manager.post_process_weights( ++ obj, request ++ ) ++ ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ + @app.post("/update_weight_version") + @auth_level(AuthLevel.ADMIN_OPTIONAL) + async def update_weight_version( +@@ -1539,6 +1564,19 @@ async def load_lora_adapter_from_tensors( + return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) + + ++@app.api_route("/update_lora_from_distributed", methods=["POST"]) ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def update_lora_from_distributed( ++ obj: Annotated[UpdateLoRAFromDistributedReqInput, Body()], request: Request ++): ++ """Load/refresh a LoRA adapter whose tensors arrive over an NCCL group.""" ++ result = await _global_state.tokenizer_manager.update_lora_from_distributed( ++ obj, request ++ ) ++ status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST ++ return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) ++ ++ + @app.api_route("/unload_lora_adapter", methods=["POST"]) + @auth_level(AuthLevel.ADMIN_OPTIONAL) + async def unload_lora_adapter( +diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +index 79136182..ef791247 100644 +--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py ++++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +@@ -2,6 +2,7 @@ from __future__ import annotations + + import contextlib + import logging ++import os + from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + + import torch +@@ -30,6 +31,7 @@ from sglang.srt.layers.attention.dsa.utils import ( + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_in_seq_split, + is_graph_dsa_split_op_surface, ++ match_head_gate_q_scale, + ) + from sglang.srt.layers.layernorm import LayerNorm, RMSNorm + from sglang.srt.layers.utils import MultiPlatformOp +@@ -294,6 +296,17 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + self.k_norm = LayerNorm( + self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 + ) ++ # NOTE: keep this override here, *after* `use_dsa_indexer_fusion` is computed ++ # above. Hoisting it earlier would let a forced `is_neox_style=False` flip that ++ # flag on and silently enable the fused indexer path. ++ env_neox_style = os.environ.get("INDEXER_ROPE_NEOX_STYLE") ++ if env_neox_style is not None: ++ if env_neox_style not in ("0", "1"): ++ raise ValueError( ++ "INDEXER_ROPE_NEOX_STYLE must be either '0' or '1' when set." ++ ) ++ is_neox_style = env_neox_style == "1" ++ + self.rotary_emb = get_rope_wrapper( + rope_head_dim, + rotary_dim=rope_head_dim, +@@ -365,6 +378,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + ): + weights = self._weights_proj_bf16_in_fp32_out(x) + weights = weights * self.n_heads**-0.5 ++ weights = match_head_gate_q_scale(weights, q_scale) + weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale + return weights + +@@ -372,6 +386,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + def _apply_q_scale_and_softmax_scale( + self, weights: torch.Tensor, q_scale: torch.Tensor + ): ++ weights = match_head_gate_q_scale(weights, q_scale) + return weights.unsqueeze(-1) * q_scale * self.softmax_scale + + @torch.compile(dynamic=True) +@@ -397,6 +412,12 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + return max_kv_len <= self.index_topk + return False + ++ def _maybe_repeat_query_heads(self, query: torch.Tensor) -> torch.Tensor: ++ if query.shape[1] < 32: ++ assert 32 % query.shape[1] == 0 ++ query = query.repeat_interleave(32 // query.shape[1], dim=1) ++ return query ++ + def _get_q_k_bf16( + self, + q_lora: torch.Tensor, +@@ -1604,6 +1625,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + query, key, weights_raw = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) ++ query = self._maybe_repeat_query_heads(query) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + self._store_index_k_cache( +@@ -1625,6 +1647,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + enable_dual_stream, + forward_batch=forward_batch, + ) ++ query = self._maybe_repeat_query_heads(query) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() +diff --git a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +index f90e7ba9..407d88d7 100644 +--- a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py ++++ b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +@@ -14,6 +14,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo + get_tc_piecewise_forward_context, + is_in_tc_piecewise_cuda_graph, + ) ++from sglang.srt.layers.attention.dsa.utils import match_head_gate_q_scale + from sglang.srt.utils import is_cuda + from sglang.srt.utils.custom_op import register_custom_op + +@@ -81,6 +82,7 @@ if _is_cuda: + ) -> torch.Tensor: + out = torch.mm(x, weight.t(), out_dtype=torch.float32) + weights = out * n_heads_inv_sqrt ++ weights = match_head_gate_q_scale(weights, q_scale) + weights = weights.unsqueeze(-1) * q_scale * softmax_scale + return weights + +diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py +index 7df042c3..89054fe6 100644 +--- a/python/sglang/srt/layers/attention/dsa/utils.py ++++ b/python/sglang/srt/layers/attention/dsa/utils.py +@@ -245,6 +245,21 @@ def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch): + return cur_cp_seq_len != 0 + + ++def match_head_gate_q_scale( ++ weights: torch.Tensor, q_scale: torch.Tensor ++) -> torch.Tensor: ++ """Broadcast head-gate weights up to ``q_scale``'s head count. ++ ++ Models whose indexer has fewer heads than the query projection (e.g. GLM-5) ++ produce ``weights`` with fewer heads than ``q_scale``; the subsequent ++ ``weights.unsqueeze(-1) * q_scale`` would otherwise broadcast incorrectly. ++ """ ++ if weights.shape[1] < q_scale.shape[1]: ++ assert q_scale.shape[1] % weights.shape[1] == 0 ++ weights = weights.repeat_interleave(q_scale.shape[1] // weights.shape[1], dim=1) ++ return weights ++ ++ + from sglang.kernels.ops.attention.dsa.cp_split import ( + dsa_cp_round_robin_split_q_seqs_kernel, + ) +diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +index 3490fc6a..b448595f 100644 +--- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py ++++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +@@ -1021,6 +1021,10 @@ class CompressedTensorsLinearMethod(LinearMethodBase): + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + ++ def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: ++ if hasattr(layer.scheme, "restore_weights_before_loading"): ++ layer.scheme.restore_weights_before_loading(layer) ++ + def create_weights( + self, + layer: torch.nn.Module, +@@ -1075,6 +1079,10 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase): + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + ++ def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: ++ if hasattr(layer.scheme, "restore_weights_before_loading"): ++ layer.scheme.restore_weights_before_loading(layer) ++ + def create_weights( + self, + layer: torch.nn.Module, +diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py +index 87b2a334..7c70342c 100644 +--- a/python/sglang/srt/lora/layers.py ++++ b/python/sglang/srt/lora/layers.py +@@ -948,6 +948,20 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): + # initializes FusedMoE with its own moe_runner for base path + super().__init__(base_layer, lora_backend) + ++ # BaseLayerWithLoRA aliases `.weight`/`.bias` so a wrapped layer's base tensors stay ++ # reachable under their un-prefixed name. FusedMoE has neither -- it exposes ++ # `w13_weight`/`w2_weight` (plus quantization scales) -- so without an alias its base ++ # tensors are only reachable as `...experts.base_layer.w13_weight`. Model `load_weights` ++ # implementations resolve the un-prefixed `experts.w13_weight` produced by ++ # `FusedMoE.make_expert_params_mapping`, so reloading base weights into a LoRA-wrapped ++ # model (RL weight sync via /update_weights_from_*) would otherwise fail. Alias every ++ # directly-registered base parameter; this covers quantized MoE (w13_weight_scale, ++ # w13_weight_packed, ...) without a hardcoded name list. `named_parameters()` is called ++ # with remove_duplicate=False on the load path, so both names remain visible. ++ for param_name, param in base_layer.named_parameters(recurse=False): ++ if not hasattr(self, param_name): ++ setattr(self, param_name, param) ++ + lora_backend.is_moe_lora = True + + self.experts_shared_outer_loras: bool = False +diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py +index 928348eb..1c5bf177 100644 +--- a/python/sglang/srt/managers/io_struct.py ++++ b/python/sglang/srt/managers/io_struct.py +@@ -1862,6 +1862,16 @@ class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): + pass + + ++class PostProcessWeightsReqInput(BaseReq, kw_only=True): ++ restore_weights_before_load: bool = False ++ post_process_quantization: bool = False ++ ++ ++class PostProcessWeightsReqOutput(BaseReq, kw_only=True): ++ success: bool ++ message: str ++ ++ + class CheckWeightsReqInput(BaseReq, kw_only=True): + action: str = "checksum" + allow_quant_error: bool = False +@@ -2166,6 +2176,45 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): + ) + + ++class UpdateLoRAFromDistributedReqInput(BaseReq, kw_only=True): ++ """Load/refresh a LoRA adapter whose tensors arrive over an NCCL process ++ group (the RLHF online-update transport), instead of via disk or a ++ serialized blob. ++ ++ Mirrors ``UpdateWeightsFromDistributedReqInput`` (metadata only; the real ++ tensors are broadcast ``src=0`` on ``group_name``) but the payload lands in ++ the ``LoRAManager`` rather than ``model.load_weights``. ``config_dict`` is ++ the HF-PEFT adapter config (same content the disk path writes to ++ ``adapter_config.json``). ++ """ ++ ++ lora_name: str ++ config_dict: Dict[str, Any] ++ names: List[str] ++ dtypes: List[str] ++ shapes: List[List[int]] ++ # Number of tensors per broadcast bucket, in ``names`` order; sums to len(names). ++ # The sender picks the boundaries and the receiver must use the identical ones or the ++ # two desync on the broadcast. ``None`` means one bucket (pre-bucketing senders). ++ bucket_sizes: Optional[List[int]] = None ++ # The NCCL group name to receive the adapter tensors on (reuses the weight-update group). ++ group_name: str = "weight_update_group" ++ # Whether to pin the LoRA adapter in memory. ++ pinned: bool = False ++ added_tokens_config: Optional[Dict[str, Any]] = None ++ # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. ++ lora_id: Optional[str] = None ++ load_format: Optional[str] = None ++ ++ def to_ref(self) -> LoRARef: ++ return LoRARef( ++ lora_id=self.lora_id, ++ lora_name=self.lora_name, ++ lora_path="__distributed__", ++ pinned=self.pinned, ++ ) ++ ++ + class LoRAUpdateOutput(BaseReq, kw_only=True): + success: bool + error_message: Optional[str] = None +@@ -2174,7 +2223,7 @@ class LoRAUpdateOutput(BaseReq, kw_only=True): + + LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = ( + LoadLoRAAdapterFromTensorsReqOutput +-) = LoRAUpdateOutput ++) = UpdateLoRAFromDistributedReqOutput = LoRAUpdateOutput + + + class BlockReqType(Enum): +diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py +index a45d122f..0334e353 100755 +--- a/python/sglang/srt/managers/schedule_batch.py ++++ b/python/sglang/srt/managers/schedule_batch.py +@@ -1103,6 +1103,7 @@ class Req(ReqDllmMixin): + self.metrics_collector = metrics_collector + if time_stats is not None: + self.time_stats = SchedulerReqTimeStats.new_from_obj(time_stats) ++ self.time_stats.disagg_mode = disagg_mode + else: + self.time_stats = SchedulerReqTimeStats(disagg_mode=disagg_mode) + self.time_stats.set_metrics_collector(metrics_collector) +@@ -2739,11 +2740,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): + + retracted_reqs = [] + first_iter = True ++ num_minimum_reqs = ( ++ 0 if server_args.disaggregation_mode == "decode" else 1 ++ ) + while first_iter or ( + not self.check_decode_mem(selected_indices=sorted_indices) + ): +- if len(sorted_indices) == 1: +- # Always keep at least one request ++ if len(sorted_indices) <= num_minimum_reqs: ++ # Unified mode keeps one request; decode disaggregation may retract all. + break + + first_iter = False +@@ -2754,7 +2758,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): + self.release_req(idx, len(sorted_indices), server_args) + + reqs_to_abort: List[Req] = [] +- if len(sorted_indices) <= 1 and not self.check_decode_mem( ++ if len(sorted_indices) <= num_minimum_reqs and not self.check_decode_mem( + selected_indices=sorted_indices + ): + # Even the last remaining request cannot fit in memory. +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index a9615a14..3fa13809 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -141,6 +141,7 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqOutput, + OpenSessionReqInput, + PauseGenerationReqInput, ++ PostProcessWeightsReqInput, + ProfileReq, + ReleaseMemoryOccupationReqInput, + RemoveExternalCorpusReqInput, +@@ -161,6 +162,8 @@ from sglang.srt.managers.io_struct import ( + TokenizedGenerateReqInput, + UnloadLoRAAdapterReqInput, + UnloadLoRAAdapterReqOutput, ++ UpdateLoRAFromDistributedReqInput, ++ UpdateLoRAFromDistributedReqOutput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -1539,6 +1542,10 @@ class Scheduler( + UpdateWeightsFromIPCReqInput, + self.weight_updater.update_weights_from_ipc, + ), ++ ( ++ PostProcessWeightsReqInput, ++ self.weight_updater.post_process_weights, ++ ), + ( + GetWeightsByNameReqInput, + self.weight_updater.get_weights_by_name, +@@ -1571,6 +1578,10 @@ class Scheduler( + LoadLoRAAdapterFromTensorsReqInput, + self.load_lora_adapter_from_tensors, + ), ++ ( ++ UpdateLoRAFromDistributedReqInput, ++ self.update_lora_from_distributed, ++ ), + (UnloadLoRAAdapterReqInput, self.unload_lora_adapter), + (PauseGenerationReqInput, self.pause_generation), + (ContinueGenerationReqInput, self.continue_generation), +@@ -4457,6 +4468,12 @@ class Scheduler( + # The request will still run one decode forward pass. + # Then we reuse all existing code to clean up the KV cache allocation. + logger.debug(f"Abort running request. {req.rid=}") ++ if self.disaggregation_mode == DisaggregationMode.PREFILL and hasattr( ++ req, "disagg_kv_sender" ++ ): ++ sender = getattr(req, "disagg_kv_sender", None) ++ if sender is not None and hasattr(sender, "abort"): ++ sender.abort() + req.to_finish = FINISH_ABORT() + + def _pause_engine(self) -> Tuple[List[Req], int]: +@@ -4669,6 +4686,14 @@ class Scheduler( + result = self.tp_worker.load_lora_adapter_from_tensors(recv_req) + return result + ++ def update_lora_from_distributed( ++ self, recv_req: UpdateLoRAFromDistributedReqInput ++ ) -> UpdateLoRAFromDistributedReqOutput: ++ """In-place loading a new lora adapter whose tensors arrive over NCCL.""" ++ ++ result = self.tp_worker.update_lora_from_distributed(recv_req) ++ return result ++ + def unload_lora_adapter( + self, recv_req: UnloadLoRAAdapterReqInput + ) -> UnloadLoRAAdapterReqOutput: +diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py +index 39eeec2e..12ba5133 100644 +--- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py ++++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py +@@ -397,7 +397,7 @@ class SchedulerProfilerManager: + if self.profiler_prefill_ct > self.profiler_target_prefill_ct: + if self.profile_in_progress: + self._stop_profile(stage=ForwardMode.EXTEND) +- elif batch.forward_mode.is_decode(): ++ elif batch.forward_mode.is_decode() or batch.forward_mode.is_prebuilt(): + if self.profiler_decode_ct == 0: + if self.profile_in_progress: + # force trace flush +diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py +index 653b28c4..d30ce9f9 100644 +--- a/python/sglang/srt/managers/scheduler_components/weight_updater.py ++++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py +@@ -28,6 +28,8 @@ from sglang.srt.managers.io_struct import ( + GetWeightsByNameReqOutput, + InitWeightsUpdateGroupReqInput, + InitWeightsUpdateGroupReqOutput, ++ PostProcessWeightsReqInput, ++ PostProcessWeightsReqOutput, + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + ResumeMemoryOccupationReqInput, +@@ -183,6 +185,19 @@ class SchedulerWeightUpdaterManager: + parameter = self.tp_worker.get_weights_by_name(recv_req) + return GetWeightsByNameReqOutput(parameter=parameter) + ++ def post_process_weights(self, recv_req: PostProcessWeightsReqInput): ++ success, message = self.tp_worker.post_process_weights(recv_req) ++ if ( ++ success ++ and self.draft_worker is not None ++ and hasattr(self.draft_worker, "post_process_weights") ++ ): ++ success, message = self.draft_worker.post_process_weights(recv_req) ++ if not success: ++ logger.error(message) ++ torch.distributed.barrier(group=self.tp_cpu_group) ++ return PostProcessWeightsReqOutput(success=success, message=message) ++ + def _assert_weight_cache_inactive(self, op: str) -> None: + """Reject freeing/restoring model weights while the CUDA IPC weight + cache is active: the weights are shared with the daemon via CUDA IPC, so +diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py +index 7cfd98b8..03ac104c 100644 +--- a/python/sglang/srt/managers/tokenizer_control_mixin.py ++++ b/python/sglang/srt/managers/tokenizer_control_mixin.py +@@ -48,6 +48,8 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqOutput, + LoRAUpdateOutput, + OpenSessionReqInput, ++ PostProcessWeightsReqInput, ++ PostProcessWeightsReqOutput, + ProfileReq, + ProfileReqOutput, + ProfileReqType, +@@ -66,6 +68,8 @@ from sglang.srt.managers.io_struct import ( + SlowDownReqOutput, + UnloadLoRAAdapterReqInput, + UnloadLoRAAdapterReqOutput, ++ UpdateLoRAFromDistributedReqInput, ++ UpdateLoRAFromDistributedReqOutput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromDistributedReqOutput, + UpdateWeightsFromIPCReqInput, +@@ -101,6 +105,7 @@ _COMMUNICATOR_SPECS = [ + ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput), + ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), + ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), ++ ("post_process_weights", PostProcessWeightsReqOutput), + ("get_weights_by_name", GetWeightsByNameReqOutput), + ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), + ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), +@@ -738,6 +743,103 @@ class TokenizerControlMixin: + error_message=str(e), + ) + ++ async def update_lora_from_distributed( ++ self: TokenizerManager, ++ obj: UpdateLoRAFromDistributedReqInput, ++ _: Optional[fastapi.Request] = None, ++ ) -> UpdateLoRAFromDistributedReqOutput: ++ self.auto_create_handle_loop() ++ ++ try: ++ if not self.server_args.enable_lora: ++ raise ValueError( ++ "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ++ ) ++ ++ assert ( ++ self.server_args.dp_size == 1 ++ ), "dp_size must be 1 for dynamic lora loading" ++ logger.info( ++ "Start update Lora adapter from distributed. Lora name=%s", ++ obj.lora_name, ++ ) ++ ++ async with self.lora_update_lock: ++ # Self-contained same-name replacement: drop the prior version of ++ # this adapter first so repeated pushes under the fixed name do not ++ # accumulate (the NCCL caller does not send a separate unload). ++ # ++ # Ask the REGISTRY ("is a version live right now?"), never lora_ref_cache ++ # ("was one ever loaded?"). Upstream's cache is append-only -- unload only ++ # touches the registry -- so a push that dies between the unload above and ++ # the receive below leaves the name in the cache but not in the registry. ++ # Keying the unload off the cache would then make every retry call ++ # unregister() on a name that is gone, which raises before the receive is ++ # ever reached: one transient failure would wedge the engine permanently. ++ if obj.lora_name in self.lora_registry.get_all_adapters(): ++ unload_result = await self._unload_lora_adapter_locked( ++ UnloadLoRAAdapterReqInput(lora_name=obj.lora_name) ++ ) ++ if not unload_result.success: ++ raise ValueError( ++ f"Error while unloading prior LoRA adapter '{obj.lora_name}': " ++ f"{unload_result.error_message}" ++ ) ++ ++ new_adapter = LoRARef( ++ lora_name=obj.lora_name, ++ lora_path="__distributed__", ++ pinned=obj.pinned, ++ ) ++ obj.lora_id = new_adapter.lora_id ++ result = (await self.update_lora_adapter_communicator(obj))[0] ++ ++ if result.success: ++ await self.lora_registry.register(new_adapter) ++ # Deliberately NOT recorded in lora_ref_cache. That cache exists so an ++ # LRU-evicted adapter can be transparently re-read from its lora_path, ++ # but these tensors arrived over NCCL and were never on disk -- ++ # "__distributed__" is a placeholder, not a real path. Caching it would ++ # send the implicit-reload path in TokenizerManager off to load a ++ # directory that does not exist; leaving it out makes an unregistered ++ # adapter surface as "never been loaded", which is the truth. ++ if self.server_args.max_loaded_loras is not None: ++ while ( ++ self.lora_registry.num_registered_loras ++ > self.server_args.max_loaded_loras ++ ): ++ lru_lora_name = await self.lora_registry.lru_lora_name( ++ exclude_pinned=True ++ ) ++ if lru_lora_name is None: ++ raise ValueError( ++ "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " ++ f"LoRA registry is: {self.lora_registry._registry}" ++ ) ++ ++ logger.info( ++ f"Unloading least recently used LoRA adapter '{lru_lora_name}' " ++ f"(current number of adapters: {self.lora_registry.num_registered_loras}, " ++ f"max allowed: {self.server_args.max_loaded_loras})" ++ ) ++ ++ unload_result = await self._unload_lora_adapter_locked( ++ UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) ++ ) ++ if not unload_result.success: ++ raise ValueError( ++ f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " ++ f"{unload_result.error_message}" ++ ) ++ del result.loaded_adapters[lru_lora_name] ++ ++ return result ++ except ValueError as e: ++ return UpdateLoRAFromDistributedReqOutput( ++ success=False, ++ error_message=str(e), ++ ) ++ + async def unload_lora_adapter( + self: TokenizerManager, + obj: UnloadLoRAAdapterReqInput, +@@ -797,6 +899,16 @@ class TokenizerControlMixin: + self.auto_create_handle_loop() + await self.resume_memory_occupation_communicator(obj) + ++ async def post_process_weights( ++ self: TokenizerManager, ++ obj: PostProcessWeightsReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ async with self.model_update_lock.writer_lock: ++ results = await self.post_process_weights_communicator(obj) ++ return FanOutCommunicator.merge_results(results) ++ + async def check_weights( + self: TokenizerManager, + obj: CheckWeightsReqInput, +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 80351f33..99e4a1c2 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -2840,27 +2840,25 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): + priority = getattr(state.obj, "priority", None) + if priority is not None: + labels["priority"] = str(priority) +- if ( +- not state.ttft_observed +- and self.disaggregation_mode != DisaggregationMode.PREFILL +- ): ++ if not state.ttft_observed: + state.ttft_observed = True + state.last_completion_tokens = completion_tokens +- self.metrics_collector.observe_time_to_first_token( +- labels, +- state.time_stats.get_first_token_latency(), +- stream=getattr(state.obj, "stream", False), +- ) ++ if self.disaggregation_mode != DisaggregationMode.PREFILL: ++ self.metrics_collector.observe_time_to_first_token( ++ labels, ++ state.time_stats.get_first_token_latency(), ++ stream=getattr(state.obj, "stream", False), ++ ) + else: + num_new_tokens = completion_tokens - state.last_completion_tokens +- if num_new_tokens: ++ if num_new_tokens > 0: + self.metrics_collector.observe_inter_token_latency( + labels, + state.time_stats.get_interval(), + num_new_tokens, + ) + state.time_stats.set_last_time() +- state.last_completion_tokens = completion_tokens ++ state.last_completion_tokens = completion_tokens + + if state.finished: + # Get detailed cache breakdown if available +@@ -3731,7 +3729,14 @@ def _get_processor_wrapper(server_args): + + + def _determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode: +- is_cross_node = server_args.dist_init_addr ++ # NOTE(relax-fix): keying "cross node" off dist_init_addr false-positives for ++ # single-node multi-GPU TP, which MUST also set dist_init_addr to bootstrap the ++ # TP process group (see relax/backends/sglang/sglang_engine.py). The false ++ # positive forces the CPU "default" transport, which makes wrap_shm_features() ++ # a no-op, so multimodal pixel tensors (multi-GB for 5-image doc pages) get ++ # pickled + gloo-broadcast per request (observed 2.2-6.5 GB, 18-33s per ++ # recv_requests). Key off nnodes, which is what actually means multi-node. ++ is_cross_node = server_args.nnodes > 1 + + if is_cross_node: + # Fallback to default CPU transport for multi-node +diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py +index 0fdf5b96..2cda3ba5 100644 +--- a/python/sglang/srt/managers/tp_worker.py ++++ b/python/sglang/srt/managers/tp_worker.py +@@ -30,8 +30,10 @@ from sglang.srt.managers.io_struct import ( + InitWeightsUpdateGroupReqInput, + LoadLoRAAdapterFromTensorsReqInput, + LoadLoRAAdapterReqInput, ++ PostProcessWeightsReqInput, + SendWeightsToRemoteInstanceReqInput, + UnloadLoRAAdapterReqInput, ++ UpdateLoRAFromDistributedReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -195,6 +197,13 @@ class BaseTpWorker(ABC): + ) + return success, message + ++ def post_process_weights(self, recv_req: PostProcessWeightsReqInput): ++ success, message = self.model_runner.post_process_weights( ++ restore_weights_before_load=recv_req.restore_weights_before_load, ++ post_process_quantization=recv_req.post_process_quantization, ++ ) ++ return success, message ++ + def _deserialize_own_rank(self, serialized_named_tensors): + """Each rank deserializes only its own payload (index ps.tp_rank); + deserializing another rank's copy would break producer-side CUDA-IPC +@@ -285,6 +294,23 @@ class BaseTpWorker(ABC): + ) + return result + ++ def update_lora_from_distributed( ++ self, recv_req: UpdateLoRAFromDistributedReqInput ++ ): ++ # Every TP worker joins the broadcast; the LoRA code slices its own shard ++ # internally (slice_lora_a_weights / slice_lora_b_weights). ++ result = self.model_runner.update_lora_from_distributed( ++ recv_req.to_ref(), ++ recv_req.names, ++ recv_req.dtypes, ++ recv_req.shapes, ++ recv_req.config_dict, ++ recv_req.group_name, ++ recv_req.added_tokens_config, ++ recv_req.bucket_sizes, ++ ) ++ return result ++ + def forward_batch_embedding(self, batch: ScheduleBatch): + forward_batch = ForwardBatch.init_new( + batch, +diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py +index 5ebac56f..e9c2d85d 100644 +--- a/python/sglang/srt/mem_cache/hiradix_cache.py ++++ b/python/sglang/srt/mem_cache/hiradix_cache.py +@@ -1137,9 +1137,7 @@ class HiRadixCache(RadixCache): + self._update_leaf_status(node) + self._update_host_leaf_status(node) + if node.parent is None: +- assert ( +- node is self.root_node +- ), f"This request holds the node from another tree" ++ break + node = node.parent + return DecLockRefResult(delta=delta) + +@@ -1245,6 +1243,7 @@ class HiRadixCache(RadixCache): + self._update_host_leaf_status(node) + # update leaf status for the parent because the node is evicted + self._update_leaf_status(node.parent) ++ self._update_host_leaf_status(node.parent) + return num_evicted + + def _evict_backuped(self, node: TreeNode): +diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py +index a9b2b8b9..27b1c48e 100644 +--- a/python/sglang/srt/mem_cache/memory_pool.py ++++ b/python/sglang/srt/mem_cache/memory_pool.py +@@ -4391,9 +4391,12 @@ class DSATokenToKVPool(MLATokenToKVPool): + def _create_index_buffers(self): + num_pages = (self.index_buf_size + self.page_size + 1) // self.page_size + with ( +- torch.cuda.use_mem_pool(self.custom_mem_pool) +- if self.custom_mem_pool +- else nullcontext() ++ ( ++ torch.cuda.use_mem_pool(self.custom_mem_pool) ++ if self.custom_mem_pool ++ else nullcontext() ++ ), ++ self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), + ): + self.index_k_with_scale_buffer = [ + torch.zeros( +diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py +index 2916f8c6..32b81c49 100644 +--- a/python/sglang/srt/mem_cache/radix_cache.py ++++ b/python/sglang/srt/mem_cache/radix_cache.py +@@ -491,6 +491,9 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): + return + + token_ids = req.get_fill_ids() ++ kv_committed_len = getattr(req, "kv_committed_len", None) ++ if kv_committed_len is not None and len(token_ids) > kv_committed_len: ++ token_ids = token_ids[:kv_committed_len] + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : len(token_ids) + ] +@@ -619,9 +622,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): + node.lock_ref -= 1 + self._update_leaf_status(node) + if node.parent is None: +- assert ( +- node is self.root_node +- ), "This request holds the node from another tree" ++ break + node = node.parent + return DecLockRefResult(delta=delta) + +diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py +index 8e957f66..1bd46d38 100644 +--- a/python/sglang/srt/model_executor/model_runner.py ++++ b/python/sglang/srt/model_executor/model_runner.py +@@ -376,9 +376,10 @@ class ModelRunner: + self.war_fastpath_read_done_event: Optional[torch.cuda.Event] = None + + # CPU offload +- set_offloader( +- create_offloader_from_server_args(server_args, dp_rank=self.ps.dp_rank) +- ) ++ if not is_draft_worker: ++ set_offloader( ++ create_offloader_from_server_args(server_args, dp_rank=self.ps.dp_rank) ++ ) + + self._weight_checker = WeightChecker(get_model=lambda: self.model, ps=self.ps) + +@@ -1183,10 +1184,130 @@ class ModelRunner: + lora_ref, tensors, config_dict, added_tokens_config + ) + ++ def update_lora_from_distributed( ++ self, ++ lora_ref: LoRARef, ++ names, ++ dtypes, ++ shapes, ++ config_dict, ++ group_name, ++ added_tokens_config=None, ++ bucket_sizes=None, ++ ): ++ """Receive a LoRA adapter's tensors over an NCCL process group and load it. ++ ++ Mirrors ``update_weights_from_distributed`` (each ``(name, dtype, shape)`` ++ is broadcast ``src=0`` on the shared weight-update group), but the ++ gathered tensors are handed to the ``LoRAManager`` instead of ++ ``model.load_weights``. The full (TP-gathered, PP-merged) adapter is ++ broadcast to every worker, which slices its own shard internally. ++ ++ ``bucket_sizes`` is the number of tensors per broadcast bucket, exactly as the ++ sender grouped them. Each bucket is staged on device, received, then moved to ++ host memory and released, so peak device memory is ONE bucket instead of the ++ whole adapter: a MoE adapter runs to several GB, and allocating that on top of ++ the KV cache OOMs mid-broadcast, which wedges every other participant until the ++ NCCL timeout. Host tensors are what ``LoRAAdapter`` keeps anyway ++ (``_process_weight`` calls ``.cpu()``). ``None`` means a single bucket. ++ """ ++ from sglang.srt.managers.io_struct import LoRAUpdateOutput ++ ++ model_update_group = self.weight_updater._model_update_group ++ assert group_name in model_update_group, ( ++ f"Group {group_name} not in {list(model_update_group.keys())}. " ++ "Please call `init_weights_update_group` first." ++ ) ++ if not bucket_sizes: ++ bucket_sizes = [len(names)] ++ assert sum(bucket_sizes) == len(names), ( ++ f"bucket_sizes sum to {sum(bucket_sizes)} but {len(names)} tensors were " ++ "announced; sender and receiver would issue different broadcasts and hang." ++ ) ++ logger.info(f"LoRA adapter loading from distributed starts: {lora_ref}.") ++ try: ++ tensors = {} ++ offset = 0 ++ for count in bucket_sizes: ++ handles = [] ++ staged = [] ++ for name, dtype, shape in zip( ++ names[offset : offset + count], ++ dtypes[offset : offset + count], ++ shapes[offset : offset + count], ++ ): ++ target_dtype = ( ++ dtype ++ if isinstance(dtype, torch.dtype) ++ else getattr(torch, dtype) ++ ) ++ weight = torch.empty(shape, dtype=target_dtype, device=self.device) ++ handles.append( ++ torch.distributed.broadcast( ++ weight, ++ src=0, ++ group=model_update_group[group_name], ++ async_op=True, ++ ) ++ ) ++ staged.append((name, weight)) ++ for handle in handles: ++ handle.wait() ++ handles.clear() ++ # Move off-device and drop the staging buffers before the next bucket is ++ # allocated, so the device high-water mark stays at one bucket. ++ for name, weight in staged: ++ tensors[name] = weight.to("cpu") ++ staged.clear() ++ offset += count ++ except Exception as e: ++ error_msg = f"Failed to receive LoRA adapter over distributed group: {e}." ++ logger.error(error_msg) ++ return LoRAUpdateOutput(success=False, error_message=error_msg) ++ ++ result = self.lora_manager.load_lora_adapter_from_tensors( ++ lora_ref, tensors, config_dict, added_tokens_config ++ ) ++ logger.info(f"LoRA adapter loading from distributed completes: {lora_ref}.") ++ return result ++ + def unload_lora_adapter(self, lora_ref: LoRARef): + """Unload a lora adapter that was previously loaded during initialization or dynamic loading.""" + return self.lora_manager.unload_lora_adapter(lora_ref) + ++ def post_process_weights( ++ self, ++ restore_weights_before_load: bool = False, ++ post_process_quantization: bool = False, ++ ): ++ """Run optional post-loading hooks, such as quantization repacking.""" ++ from sglang.srt.model_loader.loader import device_loading_context ++ ++ if self.device == "cuda": ++ target_device = torch.device("cuda", torch.cuda.current_device()) ++ else: ++ target_device = torch.device(self.device) ++ ++ if restore_weights_before_load: ++ for _, module in self.model.named_modules(): ++ quant_method = getattr(module, "quant_method", None) ++ if quant_method is not None and hasattr( ++ quant_method, "restore_weights_before_loading" ++ ): ++ with device_loading_context(module, target_device): ++ quant_method.restore_weights_before_loading(module) ++ ++ if post_process_quantization: ++ for _, module in self.model.named_modules(): ++ quant_method = getattr(module, "quant_method", None) ++ if quant_method is not None and hasattr( ++ quant_method, "process_weights_after_loading" ++ ): ++ with device_loading_context(module, target_device): ++ quant_method.process_weights_after_loading(module) ++ ++ return True, "Success" ++ + @property + def effective_max_total_num_tokens(self): + """Return the max token pool size considering hybrid swa settings.""" +@@ -1441,6 +1562,12 @@ class ModelRunner: + output.expert_distribution_metrics = recorder_outputs.get("metrics") + + no_copy_to_cpu = not get_schedule().disable_overlap_schedule ++ cuda_graph_num_tokens = None ++ decode_graph_runner = getattr(self, "decode_cuda_graph_runner", None) ++ if getattr(decode_graph_runner, "bs", None): ++ cuda_graph_num_tokens = decode_graph_runner.bs * getattr( ++ decode_graph_runner, "num_tokens_per_bs", 1 ++ ) + if ( + not self.is_draft_worker + and (experts_capturer := get_global_experts_capturer()) is not None +@@ -1448,7 +1575,7 @@ class ModelRunner: + output.routed_experts_output = experts_capturer.on_forward_end( + forward_batch=forward_batch, + can_run_graph=output.can_run_graph, +- cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None), ++ cuda_graph_batch=cuda_graph_num_tokens, + no_copy_to_cpu=no_copy_to_cpu, + ) + +@@ -1456,7 +1583,7 @@ class ModelRunner: + output.indexer_topk_output = indexer_capturer.on_forward_end( + forward_batch=forward_batch, + can_run_graph=output.can_run_graph, +- cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None), ++ cuda_graph_batch=cuda_graph_num_tokens, + no_copy_to_cpu=no_copy_to_cpu, + ) + +diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py +index 08c47b99..06a222fe 100644 +--- a/python/sglang/srt/models/qwen3_5.py ++++ b/python/sglang/srt/models/qwen3_5.py +@@ -1125,12 +1125,12 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + forward_batch: ForwardBatch, + ) -> torch.Tensor: + """Full attention forward pass.""" +- if _is_cuda and self.attn_output_gate: ++ if _is_cuda and self.attn_output_gate and positions.ndim == 1: + q, k, v, gate = self.forward_prepare_cuda_fused( + positions=positions, + hidden_states=hidden_states, + ) +- elif (_is_hip or _is_xpu or _is_cpu) and self.attn_output_gate: ++ elif (_is_cuda or _is_hip or _is_xpu or _is_cpu) and self.attn_output_gate: + q, k, v, gate = self.forward_prepare_fused_gate( + positions=positions, + hidden_states=hidden_states, +diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py +index 83697f14..aae8b5c8 100644 +--- a/python/sglang/srt/models/qwen3_vl.py ++++ b/python/sglang/srt/models/qwen3_vl.py +@@ -1171,9 +1171,14 @@ class Qwen3LLMModel(Qwen3Model): + # To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack + # The order matters because addition with different tensors is not associative in practice. + # Deepstack for prev_layer is applied at the start of current layer via post_residual_addition. +- deepstack_embeds = self.get_deepstack_embeds( +- layer_idx - 1, input_deepstack_embeds +- ) ++ deepstack_embeds = None ++ if input_deepstack_embeds is not None: ++ prev_layer_idx = layer_idx - 1 ++ if prev_layer_idx in self.deepstack_embed_to_decoder_layer: ++ sep = self.hidden_size * prev_layer_idx ++ deepstack_embeds = input_deepstack_embeds[ ++ :, sep : sep + self.hidden_size ++ ] + hidden_states, residual = layer( + positions, + hidden_states, +diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py +index db684259..17d2cb69 100644 +--- a/python/sglang/srt/multimodal/processors/glm4v.py ++++ b/python/sglang/srt/multimodal/processors/glm4v.py +@@ -1,7 +1,13 @@ + from typing import List, Union + ++import torch ++ + from sglang.srt.layers.rotary_embedding import MRotaryEmbedding +-from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput ++from sglang.srt.managers.schedule_batch import ( ++ Modality, ++ MultimodalDataItem, ++ MultimodalProcessorOutput, ++) + from sglang.srt.models.glm4v import Glm4vForConditionalGeneration + from sglang.srt.models.glm4v_moe import Glm4vMoeForConditionalGeneration + from sglang.srt.multimodal.processors.base_processor import ( +@@ -46,6 +52,8 @@ class Glm4vImageProcessor(SGLangBaseProcessor): + self.IMAGE_END_TOKEN_ID = hf_config.image_end_token_id + self.VIDEO_START_TOKEN_ID = hf_config.video_start_token_id + self.VIDEO_END_TOKEN_ID = hf_config.video_end_token_id ++ self.IM_START_TOKEN_ID = self.IMAGE_START_TOKEN_ID ++ self.IM_END_TOKEN_ID = self.IMAGE_END_TOKEN_ID + + # Vision config + self.IMAGE_FACTOR = 28 +@@ -60,6 +68,39 @@ class Glm4vImageProcessor(SGLangBaseProcessor): + video_token_id=self.IM_TOKEN_ID, + ).build(_processor) + ++ def get_mm_data(self, prompt, embeddings, img_grid_thw): ++ input_ids, offsets, _ = self.build_input_ids(prompt, img_grid_thw=img_grid_thw) ++ image_embeddings = ( ++ embeddings.get(Modality.IMAGE, embeddings) ++ if isinstance(embeddings, dict) ++ else embeddings ++ ) ++ mm_items = [ ++ MultimodalDataItem( ++ modality=Modality.IMAGE, ++ offsets=offsets, ++ precomputed_embeddings=image_embeddings, ++ ) ++ ] ++ ++ mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index_glm4v( ++ input_ids=torch.tensor(input_ids, dtype=torch.long).unsqueeze(0), ++ hf_config=self.hf_config, ++ image_grid_thw=img_grid_thw, ++ video_grid_thw=None, ++ attention_mask=None, ++ ) ++ ++ return MultimodalProcessorOutput( ++ input_ids=input_ids, ++ mm_items=mm_items, ++ im_start_id=self.IM_START_TOKEN_ID, ++ im_end_id=self.IM_END_TOKEN_ID, ++ im_token_id=self.IM_TOKEN_ID, ++ mrope_positions=mrope_positions.squeeze(1), ++ mrope_position_delta=mrope_position_delta, ++ ) ++ + def compute_mrope_positions(self, input_ids, mm_items): + image_grid_thw = None + video_grid_thw = None +diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py +index b4e37c7..eb14b7d 100644 +--- a/python/sglang/srt/multimodal/processors/base_processor.py ++++ b/python/sglang/srt/multimodal/processors/base_processor.py +@@ -1066,6 +1066,20 @@ class BaseMultimodalProcessor(ABC): + len(audios), + ) + ++ # Decoded input_ids can contain image tokens already expanded by the ++ # client processor. Raw images are processed again below, so restore ++ # one placeholder per image, as in legacy_load_mm_data. Precomputed ++ # inputs must retain their expanded token layout. ++ if ( ++ images ++ and all(not isinstance(image, dict) for image in images) ++ and isinstance(multimodal_tokens.image_token, str) ++ and multimodal_tokens.image_token_regex is not None ++ ): ++ prompt_str = multimodal_tokens.image_token_regex.sub( ++ lambda _: multimodal_tokens.image_token, prompt_str ++ ) ++ + return BaseMultiModalProcessorOutput( + images=images, + audios=audios, +diff --git a/python/sglang/srt/multimodal/processors/executor.py b/python/sglang/srt/multimodal/processors/executor.py +index 43da2ea..9fe011d 100644 +--- a/python/sglang/srt/multimodal/processors/executor.py ++++ b/python/sglang/srt/multimodal/processors/executor.py +@@ -46,7 +46,12 @@ class MultimodalProcessorExecutor: + else copy.deepcopy(self._processor) + ) + self._worker_state.processor = processor +- return function(*args, processor=processor, **kwargs) ++ try: ++ return function(*args, processor=processor, **kwargs) ++ except StopIteration as exc: ++ # asyncio cannot transfer StopIteration to its Future. Leaving it ++ # on the concurrent Future would keep the awaiting request pending. ++ raise ValueError("Multimodal processor exhausted an input iterator") from exc + + def shutdown(self) -> None: + self._executor.shutdown() +diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py +index 3bcb6b36..9b971098 100644 +--- a/python/sglang/srt/observability/req_time_stats.py ++++ b/python/sglang/srt/observability/req_time_stats.py +@@ -346,7 +346,7 @@ class ReqTimeStatsBase: + state["trace_ctx"] = TraceNullContext() + + for key in state.keys(): +- if key.endswith("time"): ++ if key.endswith("time") and state[key] > 0.0: + state[key] = convert_time_cross_thread( + state[key], + state["diff_realtime_monotonic"], +@@ -632,9 +632,19 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + + state = { + "has_timing_data": True, ++ "enable_metrics": self.enable_metrics, ++ "disagg_mode": self.disagg_mode, + "wait_queue_entry_time": self.wait_queue_entry_time, + "forward_entry_time": self.forward_entry_time, + "prefill_finished_time": self.prefill_finished_time, ++ "completion_time": self.completion_time, ++ "prefill_bootstrap_queue_entry_time": self.prefill_bootstrap_queue_entry_time, ++ "prefill_transfer_queue_entry_time": self.prefill_transfer_queue_entry_time, ++ "decode_prealloc_queue_entry_time": self.decode_prealloc_queue_entry_time, ++ "decode_transfer_queue_entry_time": self.decode_transfer_queue_entry_time, ++ "bootstrap_done_time": self.bootstrap_done_time, ++ "transfer_speed_gb_s": self.transfer_speed_gb_s, ++ "transfer_total_mb": self.transfer_total_mb, + "diff_realtime_monotonic": global_diff_realtime_monotonic, + } + return state +@@ -1149,6 +1159,13 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + + def convert_to_output_meta_info(self): + meta_data = {} ++ ++ def add_duration(key: str, start: float, end: float): ++ if start > 0.0 and end > 0.0: ++ duration = end - start ++ if duration >= 0.0: ++ meta_data[key] = duration ++ + if self.forward_entry_time > 0.0: + meta_data["forward_entry_time"] = convert_time_to_realtime( + self.forward_entry_time +@@ -1162,6 +1179,66 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + "queue_time": self.get_queueing_time(), + } + ) ++ if self.disagg_mode == DisaggregationMode.PREFILL: ++ add_duration( ++ "pd_prefill_bootstrap_queue_duration", ++ self.prefill_bootstrap_queue_entry_time, ++ self.wait_queue_entry_time, ++ ) ++ add_duration( ++ "pd_prefill_bootstrap_duration", ++ self.prefill_bootstrap_queue_entry_time, ++ self.bootstrap_done_time, ++ ) ++ add_duration( ++ "pd_prefill_alloc_wait_duration", ++ self.bootstrap_done_time, ++ self.wait_queue_entry_time, ++ ) ++ add_duration( ++ "pd_prefill_forward_duration", ++ self.forward_entry_time, ++ self.completion_time, ++ ) ++ add_duration( ++ "pd_prefill_transfer_queue_duration", ++ self.prefill_transfer_queue_entry_time, ++ self.completion_time, ++ ) ++ if self.transfer_speed_gb_s > 0.0: ++ meta_data["pd_transfer_speed_gb_s"] = self.transfer_speed_gb_s ++ if self.transfer_total_mb > 0.0: ++ meta_data["pd_transfer_total_mb"] = self.transfer_total_mb ++ elif self.disagg_mode == DisaggregationMode.DECODE: ++ add_duration( ++ "pd_decode_prealloc_duration", ++ self.decode_prealloc_queue_entry_time, ++ self.decode_transfer_queue_entry_time, ++ ) ++ add_duration( ++ "pd_decode_bootstrap_duration", ++ self.decode_prealloc_queue_entry_time, ++ self.bootstrap_done_time, ++ ) ++ add_duration( ++ "pd_decode_alloc_wait_duration", ++ self.bootstrap_done_time, ++ self.decode_transfer_queue_entry_time, ++ ) ++ add_duration( ++ "pd_decode_transfer_duration", ++ self.decode_transfer_queue_entry_time, ++ self.wait_queue_entry_time, ++ ) ++ add_duration( ++ "pd_decode_forward_duration", ++ self.forward_entry_time, ++ self.completion_time, ++ ) ++ if self.transfer_speed_gb_s > 0.0: ++ meta_data["pd_transfer_speed_gb_s"] = self.transfer_speed_gb_s ++ if self.transfer_total_mb > 0.0: ++ meta_data["pd_transfer_total_mb"] = self.transfer_total_mb + return meta_data + + def format_duration(self, duration: float) -> str: +diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +index 13371f43..1bd91856 100644 +--- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py ++++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +@@ -566,8 +566,10 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): + forward_batch.seq_lens, + forward_batch.out_cache_loc, + forward_batch.positions, +- forward_batch.spec_info.topk_p, +- forward_batch.spec_info.topk_index, ++ forward_batch.spec_info.topk_p.clamp(0.0, 1.0), ++ forward_batch.spec_info.topk_index.clamp( ++ 0, self.model_runner.model_config.vocab_size - 1 ++ ), + forward_batch.req_pool_indices, + ] + if buffers.rids_int is not None and forward_batch.rids_int is not None: +diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py +index 6636d06a..53142e01 100644 +--- a/python/sglang/srt/utils/common.py ++++ b/python/sglang/srt/utils/common.py +@@ -2834,6 +2834,9 @@ class SafeUnpickler(pickle.Unpickler): + "sglang.srt.utils.", + "sglang.srt.disaggregation.", + "sglang.srt.managers.", ++ "slime.", ++ # --- Relax --- ++ "Relax.", + "torch_npu.", + } + diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index c1afd6b95..3e5ef1c0f 100644 --- a/docs/.vitepress/config.mts +++ b/docs/.vitepress/config.mts @@ -277,13 +277,15 @@ export default defineConfig({ { text: 'Metrics Service', link: '/en/guide/metrics-service-detailed' }, { text: 'Notification System', link: '/en/guide/notification-system' }, { text: 'Update Weights Pipeline', link: '/en/guide/update-weights-pipeline' }, - { text: 'Low-Rank Adaptation (LoRA) Training', link: '/en/guide/low-rank-adaptation-training' } + { text: 'Low-Rank Adaptation (LoRA) Training', link: '/en/guide/low-rank-adaptation-training' }, + { text: 'Diffusion Generative RL', link: '/en/guide/diffusion-generative-rl' } ] }, { text: 'Best Practices', items: [ { text: 'Performance Tuning', link: '/en/guide/performance-tuning' }, + { text: 'Compiler Cache Reuse', link: '/en/guide/compiler-cache' }, { text: 'Accelerated S3 Model Loading', link: '/en/guide/s3-model-loading' }, { text: 'OOM Troubleshooting', link: '/en/guide/oom-troubleshooting' }, { text: 'External Model Integration', link: '/en/guide/external-model-integration' } @@ -393,13 +395,15 @@ export default defineConfig({ { text: 'Metrics 服务', link: '/zh/guide/metrics-service-detailed' }, { text: '通知系统', link: '/zh/guide/notification-system' }, { text: '权重更新流水线优化', link: '/zh/guide/update-weights-pipeline' }, - { text: '低秩适配（LoRA）训练', link: '/zh/guide/low-rank-adaptation-training' } + { text: '低秩适配（LoRA）训练', link: '/zh/guide/low-rank-adaptation-training' }, + { text: '扩散生成式 RL', link: '/zh/guide/diffusion-generative-rl' } ] }, { text: '最佳实践', items: [ { text: '性能调优', link: '/zh/guide/performance-tuning' }, + { text: '编译缓存复用', link: '/zh/guide/compiler-cache' }, { text: 'S3 模型加载加速', link: '/zh/guide/s3-model-loading' }, { text: 'OOM 排查', link: '/zh/guide/oom-troubleshooting' }, { text: '外部模型接入', link: '/zh/guide/external-model-integration' } diff --git a/docs/draft/multimodal-gen-rl-design.md b/docs/draft/multimodal-gen-rl-design.md new file mode 100644 index 000000000..9237d0fe6 --- /dev/null +++ b/docs/draft/multimodal-gen-rl-design.md @@ -0,0 +1,1167 @@ +# 多模生成 RL 适配设计 + +> **当前发布范围：Qwen-Image T2I。** FlowGRPO + FSDP2（全参数或 LoRA）+ SGLang Native Diffusion + 同步 colocate。已验证配方是 LoRA adapter-sync：100 轮完整跑通并对齐 Qwen-Image LoRA 参考曲线（最终 eval 0.8643，对照 0.8620）。 +> +> **设计上的多模型范围（Qwen-Image-Edit I2I、WAN 2.2 T2V/I2V/V2V、LTX-2.3 T2AV）已下线。** 完整代码只保留在**本地** git 分支 `backup/diffusion-generative-rl-full`（commit `f3e7e9e2`、`f7e30890`）。 +> +> ::: danger 备份分支尚未推送 +> 该分支和这两个 commit **只存在于本地仓库，`origin` 上没有**，所以对任何其他读者来说这个指针目前是无效的。若要让这条扩展参考真正可用，必须先 `git push origin backup/diffusion-generative-rl-full`；在那之前，请把本文中所有「已退役到备份分支」理解为「代码已从工作分支删除，只在某台机器的本地分支上」。 +> ::: +> +> 设计基线：Relax `e59cd7288995b13d6fac8cfbae2b638e45fb29b7`，SGLang `4ad418d2c3d43cb3c699bc9419d32673b1fca7d8`。初版日期 2026-07-16；本次按仓库实现回写，日期 2026-08-10。 +> +> 阅读约定：本文既是设计文档也是实现记录。**「已实现」**表示工作分支上有对应代码；**「已退役」**表示代码只在上述本地备份分支；**「设计保留、未实现」**表示只有设计、没有代码。 + +## 0. 现状总结 + +### 0.1 实际发布的东西 + +| 维度 | 现状 | +|---|---| +| 任务 | 只有 `t2i`（Qwen-Image）。`GENERATION_TASKS = ("t2i",)` —— 词表已随代码收敛，自带 adapter 的调用方必须同时把任务名加进来 | +| 算法 | FlowGRPO（组内中心化 advantage + clipped policy loss），`--advantage-estimator grpo` | +| 训练后端 | `--train-backend fsdp` → `FSDPTrainRayActor`，FSDP2，纯 DP（无 TP/PP/CP）；`--fsdp-trainable-mode` 支持 `full` 与 `lora` | +| Rollout 后端 | `SGLangNativeGenerationEngine`（SGLang 主线 native diffusion + Relax 静态 patch 契约检查） | +| Reward | PickScore（`relax/engine/rewards/pickscore.py`） | +| 部署模式 | 仅同步 colocate（`--colocate`），fully-async / hybrid 在预检直接拒绝 | +| 权重同步 | 每步同步。full/merge：完整 transformer，DTensor `full_tensor()` → 分桶 → CUDA-IPC → engine commit，带 count/bytes/checksum 校验与版本回滚；adapter 模式只发 LoRA 张量 | +| Checkpoint | DCP sharded（LoRA 只存 adapter）+ `COMMITTED` 协议 + 保留轮转 + 离线 HF/diffusers 与 HF-PEFT 导出 | +| 验证入口 | `scripts/training/diffusion/run-qwen-image-t2i-lora-8xgpu.sh`（默认 adapter sync，唯一复现参考曲线的配方）；`run-qwen-image-t2i-8xgpu.sh` 是全参参考 | + +### 0.2 已退役到 `backup/diffusion-generative-rl-full` 的东西 + +```text +relax/models/wan_video/{__init__.py,adapter.py} # WAN 2.2 T2V/I2V/V2V +relax/models/ltx_av/{__init__.py,adapter.py} # LTX-2.3 T2AV +relax/engine/rewards/editreward.py # EditReward (I2I) +relax/engine/rewards/video_reward.py # VideoPickScore / condition consistency / VideoCLIPDelta +relax/engine/rewards/t2av_reward.py # CLAP + 复合 reward +examples/diffusion/qwen_image_edit/i2i_full.yaml +examples/diffusion/wan22/{t2v,i2v,v2v}_full.yaml +examples/diffusion/ltx23/t2av_full.yaml +tests/models/{wan_video,ltx_av}/test_adapter.py +tests/engine/rewards/test_{editreward,video_reward,t2av_reward}.py +``` + +退役理由：这些路径没有 GPU 端到端验证，保留在主干上会让「配置得出来但跑不通」的表面积远大于实际可用面。 + +**这次退役不是纯删文件，抽象也一起收敛了**（与初版设计的最大差异，见 §4.1）：多 track（video + audio）支撑已经删除 —— `policy_tracks`、`track_weights`、`combine_track_logp`、`TrackLogp` 都不在了，`replay_transition` 直接返回 `torch.Tensor`。适配器协议同时去掉了 `encode_conditions` 与 `freeze_non_policy_modules`（冻结在 `load_train_model` 里一次做完）。因此从备份分支回接 T2AV 需要重新引入多 track 的 log-prob 合成，不再是「加一个 adapter 就行」。 + +`QwenImageAdapter.supported_tasks` 现在也只有 `("t2i",)`：i2i 分支曾经被声明为支持，但 `replay_transition` 会静默忽略源图，也就是在训练一个无条件速度场 —— 这是把词表和实际 replay 路径对齐的直接原因。 + +### 0.3 与初版设计相比的语义变化 + +1. **`--num-updates-per-batch` 的语义是切分，不是复用**（§6.4）。`_plan_micro_batch_updates` 把本 rank 的训练切片划成 N 份互不相交、样本数相等的更新，每个样本只参与一次 optimizer step。因此真正的 step 边界是 `global_batch_size / num_updates_per_batch`。初版设计里「每个 prompt group 一次 step」和中间那版「`--global-batch-size` 才是 step 边界」都已不成立，`--fsdp-optimizer-step-per-group` 也已删除。 +2. **新增 LoRA**（§10.1）：`--fsdp-trainable-mode lora` + 共享的 `--lora-*` 参数，merge / adapter 两条 rollout 同步路径。 +3. **新增按 optimizer step 计数的 LR schedule**（§10.1）：`--fsdp-lr-scheduler {constant,linear,cosine}`，默认 `constant`。 +4. **`sde_indices` 可从 `num_sde_steps` + `sde_timestep_fraction` 推导，且默认被每轮重抽样覆盖**（§5.5 / §13）。 +5. **评测、checkpoint 轮转、权重同步校验、artifact 保留、profiler 从「计划中」变成「已实现」**。 +6. **不支持的 flag 从静默无效变成预检报错**（§13.1）。 +7. **`sde_type="sde"` 下训练 SDE step 0 被硬拒绝**（§6.1 / §13.1）。 + +### 0.4 已知限制（诚实清单） + +| 限制 | 说明 | +|---|---| +| 只有同步 colocate | `fully_async` / `hybrid` / 非 colocate 在 `validate_generative_config` 被拒；`update_weights_fully_async` 抛 `NotImplementedError` | +| 训练侧纯 DP | `_get_parallel_config` 固定 `tp_size=1, pp_size=1`；rollout 侧 TP 由 `--rollout-num-gpus-per-engine` 支持 | +| replay 的 DP 切分不完整 | rank 0 读整个 TQ partition 再 broadcast 数值行（`_read_train_partition` 里的 `TODO(agent)`）。组内候选按 `[dp_rank::dp_world]` 切分（`hydrate_micro_batches`），但 `group_size % dp_world != 0` 时退化为每 rank 全量重放（结果正确，算力冗余） | +| 无训练数据 dump / debug replay | `--save-debug-train-data` 只告警；`--load-debug-rollout-data` 直接拒绝。TQ 行是数值索引 + 磁盘 sidecar，不是 token 序列 | +| 无 MFU / TFLOPs | 共享 `FlopsCounter` 是 LLM 架构专用的，DiT 没有解析模型；改用 transitions/samples per second（§14） | +| 不支持 CFG | `guidance_scale != 1.0` 在 `build_rollout_request` 直接报错：actor 只回放一次正向条件前向，双前向 CFG replay 未实现 | +| 依赖 SGLang 静态 patch | `docker/patch/sglang/v0.5.15.post1.patch` 必须应用（内存态权重同步、`/set_lora_from_tensor`、offload 端点，以及 rollout/replay 一致性改动）；engine 启动前做只读契约检查（§9） | +| KL-to-ref 未实现 | 无 reference model、无 KL 项；`kl_coef` / `use_kl_loss` 非零直接报错 | +| resume 需要显式 `--load` | `_maybe_resume` 读的是 `--load` 而不是 `--save`，两个启动脚本都没设置它 | +| reward 只有进程内或 remote | `post_process` 跑在 rollout worker 里，没有 controller 创建的 placement group。`colocate` / `cpu` 都是进程内评分（区别只在允不允许用 CUDA），`remote` 走 HTTP；旧的 `dedicated` 只是 `colocate` 的同义词并额外打一条假告警，已删除 | + +## 1. 目标方案 + +### 1.1 任务矩阵 + +| 任务 ID | 输入 | 输出 | 模型 | Rollout | 在线 reward | 状态 | +|---|---|---|---|---|---|---| +| `t2i` | text | image | `Qwen/Qwen-Image` | SGLang Native Diffusion | PickScore | **已实现并验证**（`GENERATION_TASKS` 中唯一项） | +| `i2i` / `t2v` / `i2v` / `v2v` / `t2av` | — | — | Qwen-Image-Edit、WAN 2.2、LTX-2.3 | — | EditReward / VideoPickScore / CLAP 等 | 已退役（见 §0.2 与 §4.3）。任务 ID 也已从 `GENERATION_TASKS` 移除 | + +每个训练任务独立启动，一个进程组只持有一个模型族和一个任务配置。不做跨模型族的混合 batch 或联合 optimizer。 + +### 1.2 统一训练口径 + +1. 训练完整 diffusion/flow transformer；可选用 LoRA 只训练注入的 adapter（`--fsdp-trainable-mode lora`）。 +2. VAE、text encoder、image encoder 保持冻结，且不在 FSDP actor 的模块里 —— 它们只存在于 rollout engine。 +3. 训练后端固定为 PyTorch FSDP2，参数与 optimizer 使用 DCP sharded checkpoint（LoRA 只存 adapter）。 +4. 算法固定为 FlowGRPO，rollout 保存被选中的 SDE transition，actor replay 产生 `old_logp`。 +5. actor 与 rollout 使用同步 colocate。fully async 只保留 full-weight transport 的接口形状，未实现。 +6. GPU 数量由任务 profile 决定，没有全局 8 卡上限。跨节点 FSDP 与多 GPU rollout engine 都是合法配置。 +7. SGLang 主仓库的 native diffusion 是统一推理后端；SGLang-Omni 不进入这些 diffusion 任务的运行路径。 + +### 1.3 架构边界 + +运行时代码按 Relax 已有 ownership 放置，不创建 `relax/diffusion/` 顶层领域包： + +- FSDP 状态、checkpoint 和完整权重导出进入 `relax/backends/fsdp/`。 +- SGLang DiffGenerator 生命周期进入 `relax/backends/sglang/`。 +- 通用生成模型协议和 FlowGRPO 数学进入 `relax/models/`。 +- 各模型族的 condition、trajectory geometry 和参数边界进入各自模型目录（现在只剩 `relax/models/qwen_image/`）。 +- rollout 编排进入 `relax/engine/rollout/`。 +- scorer 实现进入 `relax/engine/rewards/`，Ray 资源生命周期进入 `relax/distributed/ray/`。 + +控制面继续使用现有 `Controller`、`Actor`、`Rollout`、`RolloutManager` 和 `grpo` 注册键。数据面继续使用现有 `RolloutDataSource`、`Sample`、`TransferQueue` 与 `GRPOGroupNSampler`。 + +## 2. Relax 能力复用 + +| 现有能力 | 使用方式 | 落地情况 | +|---|---|---| +| `Sample.multimodal_inputs` | 承载条件 image/video | 已实现（t2i 不使用） | +| `MultimodalTypes` | 复用 image/video/audio 类型与 placeholder | 生成输出走 artifact manifest，未新增 Sample 字段 | +| `RolloutDataSource` | 读取统一 JSONL，按 `n_samples_per_prompt` 展开 group | 已实现；额外改动：`hf_checkpoint` 为空时不加载 tokenizer/processor | +| `_shallow_copy_sample()` | 同组候选共享只读条件媒体 | 复用现状 | +| `--multimodal-keys` | 映射数据列到 image/video/audio | 由任务 YAML 配置；t2i 为 `null` | +| `RolloutManager` | 生命周期、健康检查、offload/onload、engine recovery | 已实现，engine class 由 `_resolve_rollout_engine_class` dotpath 解析 | +| `--rollout-function-path` | 指向统一 native generation driver | 已实现（`generate_rollout`，`evaluation=True` 分流到 `evaluate_rollout`） | +| `custom_reward_post_process_path` | 完整 group 生成后批量评分 | 已实现（`relax.engine.rewards.generative.post_process`） | +| `custom_convert_samples_to_train_data_path` | 生成轻量 diffusion TQ row | 已实现（`relax/utils/utils.py` 执行该 hook） | +| `TransferQueue` | 每个候选一条数值 row，partition 仍是训练屏障 | 已实现；大 trajectory 走磁盘 sidecar | +| `GRPOGroupNSampler` | 保证同 prompt 候选组完整分配 | 复用现状 | +| 共享 `--lora-*` flag | rank / alpha / dropout / target modules / merge-adapter 模式 | 已实现；注入与同步是 diffusers 形态的独立实现（`backends/fsdp/lora.py`），只在导出时复用 `megatron_peft_utils.write_hf_peft_adapter` 一类与后端无关的 helper | +| `UpdateWeightFromTensor` | pause、分桶、Ray IPC、engine fan-out、版本校验 | FSDP actor 独立实现 `_run_weight_transaction`，mirror 其 per-rank→per-engine 映射（不复用 Megatron iterator） | +| DCS coordinator | async 分支的 topology 与版本协调 | **设计保留、未实现**：`weight_update.py` 没有 FSDP comm backend factory，`checkpoint_service/client/engine.py` 未改动 | +| `train_dump_utils` | rollout JSONL 与 debug dump | **设计保留、未实现**：该文件未改动，未透传 artifact/trajectory manifest 摘要 | +| DCP/轮转目录约定 | `iter_0000001`、latest marker 与保留策略 | 已实现；`rotate_ckpt` 新增 `save_dir` 参数以匹配 `//iter_*` 布局 | +| `TrainProfiler` | torch / memory profiler | 已实现（`--use-pytorch-profiler`、`--record-memory-history` 等对 FSDP actor 生效） | + +以下核心对象保持原样（未因本方案改动其逻辑）： + +```text +relax/core/registry.py +relax/core/controller.py +relax/core/service.py +relax/components/actor.py +relax/components/rollout.py +relax/utils/types.py +relax/utils/multimodal/* +``` + +同步 colocate 使用 `ROLES_COLOCATE`。FlowGRPO advantage 在 rollout 侧的 reward post-process 内计算、replay 与 loss 在 FSDP actor 内完成，不启动独立 Advantages Serve。 + +> **修正**：初版写「不新增或修改 `relax/components/advantages.py`」。实际做了一处**非语义**改动 —— 把 `from megatron.core import mpu` 改成函数内延迟导入，使控制面在无 megatron 的 diffusion 镜像里可以 import。同类改动还有 `relax/engine/sft/eval/runner.py`、`relax/utils/data/stream_dataloader.py`、`relax/utils/rocm_checkpoint_writer.py`、`relax/distributed/checkpoint_service/{utils.py,backends/device_direct.py}`、`relax/models/__init__.py`（Qwen-Omni import 容错）。这些都是「megatron 变可选依赖」的兼容改动，不改变 token RL 行为。 + +## 3. 总体架构 + +```mermaid +flowchart LR + C["Controller / grpo"] --> A["components.Actor"] + C --> R["components.Rollout"] + A --> F["FSDPTrainRayActor"] + R --> M["existing RolloutManager"] + M --> E["SGLangNativeGenerationEngine"] + M --> D["existing RolloutDataSource"] + D --> S["standard Sample groups"] + E --> Q["QwenImageAdapter"] + E --> T["trajectory sidecars"] + E --> O["image artifacts"] + O --> RM["generative reward post_process"] + RM --> P["PickScoreScorer"] + RM --> X["existing TransferQueue"] + X --> F + F --> U["full-weight bucket sync (CUDA IPC)"] + U --> E +``` + +> 备份分支上同一张图还有 `WAN adapter` / `LTX adapter` 两个 engine 下游分支和 `EditReward/Video/CLAP` 三类 scorer；现在只剩 Qwen 一条。 + +### 3.1 Colocate 资源模型 + +actor 与 rollout 共享同一 placement group。每个任务满足： + +```text +actor_world_size == rollout_total_gpus +rollout_total_gpus % rollout_num_gpus_per_engine == 0 +num_rollout_engines = rollout_total_gpus / rollout_num_gpus_per_engine +``` + +Reward 只有两种落点：`post_process` 所在的 rollout worker 进程内（`colocate` 允许用 CUDA，`cpu` 禁用 CUDA），或外部 HTTP 服务（`remote`）。没有独立的 reward placement group —— `GenerativeRewardManager` 现在就是 remote 的 HTTP 代理，`LocalGenerativeRewardManager` 是进程内实现。两个启动脚本都用 `--reward-runtime colocate`：PickScore 约 5.6 GB，而它运行时 actor 已经 offload。 + +任务 profile 只提供容量规划起点，不构成固定上限： + +| Profile | 建议起始 GPU | Rollout TP | Reward | 状态 | +|---|---:|---:|---|---| +| Qwen T2I LoRA | 8 x 96GB | 1 | colocate PickScore | **已验证**：`rollout_batch_size=32`、`n_samples_per_prompt=8`、384x384、12 步、训练 3 个 SDE step、rank 64 adapter | +| Qwen T2I 全参 | 8 x 96GB | 1 | colocate PickScore | 全参参考：`rollout_batch_size=8`、`n_samples_per_prompt=8`、384x384（曲线未复现） | +| Qwen T2I（YAML profile） | 8 x 96GB | 1 | colocate PickScore | `examples/diffusion/qwen_image/t2i_{full,lora}.yaml`：48 prompts x 16 / 24 candidates（未在本环境验证，作为参考片段保留） | +| Qwen Edit / WAN / LTX | — | — | — | 已退役 | + +> **Rollout TP 分片权重同步已实现。** engine 跨 `rollout_num_gpus_per_engine` 张物理卡，与之同卡的 `tp` 个 FSDP rank 组成一个 Gloo gather group，向 engine 发送 `tp` 个 CUDA-IPC blob（worker `i` open physical GPU `base+i` 上的 handle 后内部再分片）—— mirror megatron `UpdateWeightFromTensor` 的 colocate IPC 路径（`_build_ipc_gather_topology` / `_run_weight_transaction`）。拓扑按 CUDA device UUID 匹配，不按 GPU 序号，因为 placement group 会重排。唯一硬约束是 `actor_world % rollout_num_gpus_per_engine == 0`（`validate_generative_config` fail-fast，运行时再校验 `world == sum(engine_gpu_counts)`）。`tp == 1` 退化为每 rank 一个单卡 engine。 + +OOM 时增加 actor/rollout world size、提高 rollout TP、减少 `n_samples_per_prompt`、缩小训练的 SDE step 子集，或改用 LoRA。Validator 不执行模式降级，不把 full-FT 自动切成 LoRA。 + +### 3.2 Owner 状态机 + +```text +TRAIN + -> WEIGHT_SYNC + -> ROLLOUT + -> REWARD + -> TRANSFER_COMMIT + -> TRAIN +``` + +- `TRAIN`：只有 FSDP parameter/optimizer shard 在计算设备。 +- `WEIGHT_SYNC`：FSDP parameter shard 与 SGLang weight receiver 同驻，不运行 forward/backward/generate。此时 engine 只 onload transformer（`tags=[WEIGHTS]`）。 +- `ROLLOUT`：FSDP parameter 和 optimizer offload，SGLang encoder/DiT/VAE 全量 onload。 +- `REWARD`：全部 artifact 已提交；scorer 评分。 +- `TRANSFER_COMMIT`：完整 group 校验后写入 TQ partition，seal 后训练才可消费。 + +实现上 `REWARD` 与 `TRANSFER_COMMIT` 都发生在 `generate_rollout` 尾部的 `transfer_batch_to_data_system` 调用链里（reward post-process → converter → `async_put`），所以 `perf/reward_time` 覆盖的是「评分 + 转换 + TQ 落盘」三段之和。 + +## 4. 任务与模型适配 + +### 4.1 通用适配协议（已实现，`relax/models/generative.py`） + +```python +@runtime_checkable +class GenerativeModelAdapter(Protocol): + family: str + supported_tasks: tuple[str, ...] + + def load_train_model(self, config) -> torch.nn.Module: ... + def build_rollout_request(self, sample, sampling, seed) -> dict: ... + def validate_rollout_response(self, response) -> None: ... + def pack_trajectory(self, response) -> dict[str, torch.Tensor]: ... + def replay_transition(self, model, batch, step_index) -> torch.Tensor: ... + def artifact_tracks(self, response) -> list[ArtifactTrack]: ... + def weight_name_map(self, name: str) -> str: ... +``` + +另有两个**可选属性**，用 `getattr` 读取、刻意不声明为协议成员（这样加属性不会让已有 adapter 失效 —— 协议是 `runtime_checkable` 的，adapter 单测会 assert 这一点）：`lora_target_modules`（`--fsdp-trainable-mode lora` 的默认目标模块，模块名后缀）与 `lora_task_type`（PEFT task type，默认 `FEATURE_EXTRACTION`）。 + +相对初版设计，协议**收窄**了三处： + +- 去掉 `policy_tracks` 与多 track 支撑（见 §0.2 / §6.2）；`replay_transition` 直接返回 `torch.Tensor`。 +- 去掉 `freeze_non_policy_modules`：`load_train_model` 返回的模块会被直接交给 FSDP，所以「不该拿梯度的东西」必须在那里就设好 `requires_grad=False`，没有第二个冻结钩子。 +- 去掉 `encode_conditions`：条件由 engine 在 `denoising_env` 里回传，adapter 不再自己编码。 + +任务配置通过 `--model-adapter-path`（dotpath）选择实现。FSDP runtime、rollout driver、reward barrier 和 weight transport 只依赖该协议。 + +同模块还定义 `ArtifactTrack`、`FullWeightManifest`、`artifact_manifest_dict()`、`ordered_name_shape_hash()` 和 `resolve_sde_indices()`。 + +### 4.2 Qwen-Image T2I（已实现并验证） + +- policy：完整 `QwenImageTransformer2DModel`，`--fsdp-trainable-attr transformer`；LoRA 时只训练八个 attention 投影上的 adapter。 +- frozen：Qwen2.5-VL text encoder、VAE、scheduler —— 它们不在 FSDP actor 的模块里，只存在于 rollout engine。 +- condition：`prompt_embeds` 与 `prompt_embeds_mask`，由 engine 在 `denoising_env` 里回传。 +- output：PNG image track，从 `generated_output` 落盘到 `//rollout_*/group_*_sample_*.png`。 +- trajectory：packed latent（`image_x_t` / `image_x_next`）、`sigmas`、`sde_indices`、`image_grid`、seed hash。`image_grid` 由 engine 回显的 `height` / `width` 算出 —— 打包后的序列长度无法还原非正方形网格，所以这是重建 RoPE 的唯一来源。 +- reward：PickScore，组内中心化后除以该组自己的 std。 +- 硬约束：`guidance_scale` 必须为 `1.0`（`replay_transition` 只跑一次正向条件前向）。 + +### 4.3 已退役的模型族（WAN 2.2 / LTX-2.3 / Qwen-Image-Edit） + +这些设计的完整推理与代码都在备份分支上（见文首的推送提醒），本文不再复述，只留必要的回接提示： + +- **Qwen-Image-Edit I2I**：`QwenImageAdapter` 中的 i2i 分支已删除，`supported_tasks` 收敛为 `("t2i",)`。回接时要同时补上 source image 条件、`denoising_env.image_kwargs` 校验、EditReward scorer 与 aspect bucket 分组。 +- **WAN 2.2 T2V/I2V/V2V**：5D latent `[B,C,T,H,W]`、frame count 与 VAE temporal factor 的约束、MP4 track manifest、VideoPickScore / condition consistency / VideoCLIPDelta。 +- **LTX-2.3 T2AV**：video + audio 双 track、时长同步校验、CLAP 与 VideoPickScore 的加权合成。 + +::: warning 回接成本已经变高 +多 track 的通用支撑**已经删除**（`combine_track_logp`、`TrackLogp`、`track_weights` 都不在了），所以 T2AV 不再是「换个 adapter」就能回来的：需要重新引入按 track 加权合成 log-prob 的那一层。视频/音频任务的 `GENERATION_TASKS` 条目也要一并加回。 +::: + +## 5. 数据契约 + +### 5.1 统一 JSONL + +T2I（唯一在用的形态）： + +```json +{"prompt":"a tram passing through a rainy city at night","metadata":{"task":"t2i","prompt_id":"pickapic_0001"}} +``` + +条件类任务（设计保留）： + +```json +{"prompt":" replace the sky with a sunset","images":["/data/source/0001.png"],"metadata":{"task":"i2i","sample_id":"magicbrush_0001"}} +{"prompt":"