From 8e6d9f563c5455d886fc30cbcc68429b3c405485 Mon Sep 17 00:00:00 2001 From: sopyang Date: Thu, 10 Sep 2026 01:38:12 -0700 Subject: [PATCH 1/6] [megatron] feat: MXFP4 experts-only QAT under TransformerEngine FP8 GEMM The DeepSeek-V4 report describes FP4 quantization-aware training applied to MoE expert weights during post-training. This adds `qat.mode=mxfp4_experts`: a weight-only ModelOpt fake-quant onto the MXFP4 grid (E2M1 values, one E8M0 power-of-two scale per 32 contiguous K elements) scoped to routed experts via `*mlp.experts*weight_quantizer`, with `algorithm=None` because dynamic block quantization derives its scale per call and ModelOpt asserts it is never calibrated. The GEMMs keep running in real FP8 through TE's blockwise autocast rather than falling back to a simulated BF16 matmul. ModelOpt refuses to run any quantizer under `fp8_autocast` -- a guard that keeps pure simulation faithful, but here the stacking is the point (FP4 weights feeding a real FP8 GEMM, which is what the checkpoint's own mixed MXFP4/FP8 layout implies). `qat.bypass_te_fp8_assert` lifts it at runtime without patching ModelOpt. Activations are quantized once, by TE; ModelOpt's input_quantizer stays disabled so they are not quantized twice. Two adjacent fixes were required to make QAT usable on such a model: - The ModelOpt layer spec is no longer forced on providers that declare an `experimental_attention_variant` (DeepSeek-V4 "dsv4_hybrid", GLM5 "dsa", Qwen3-Next "gated_delta_net"). Those install their own heterogeneous spec in `provider_bridge()`, and the generic ModelOpt GPT spec would silently replace it and build a different model. ModelOpt swaps QuantModules by class, so the spec change is not needed for quantization to apply. - `export_qat_weights` is gated to the NVFP4 modes it actually implements. For `mxfp4_experts` the bridge already re-quantizes routed experts to MXFP4 on every weight sync, so the training forward and the rollout weights land on one grid without a second exporter. `apply_qat_to_modules` now fails fast if the experts-only config did not land where intended: it requires at least one enabled routed-expert weight quantizer and no other enabled quantizer (a mis-ordered cfg would leave every remaining linear on ModelOpt's default per-tensor INT8), and logs whether the fake-quant is the identity on the loaded weights -- it is, exactly, when the checkpoint already stores those experts on the MXFP4 grid. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: OpenAI Codex --- verl/trainer/config/engine/megatron.yaml | 5 +- verl/utils/modelopt/qat_utils.py | 52 +++++++++++++++++++ verl/utils/modelopt/quantize.py | 32 ++++++++++-- verl/workers/config/engine.py | 3 ++ .../engine/megatron/transformer_impl.py | 23 +++++++- 5 files changed, 108 insertions(+), 7 deletions(-) diff --git a/verl/trainer/config/engine/megatron.yaml b/verl/trainer/config/engine/megatron.yaml index 01c5dc51fd7..db64d2bd090 100644 --- a/verl/trainer/config/engine/megatron.yaml +++ b/verl/trainer/config/engine/megatron.yaml @@ -144,7 +144,8 @@ qat: _target_: verl.workers.config.QATEngineConfig # Whether to enable QAT enable: false - # Quantization mode: "w4a16" (weight-only FP4) or "w4a4" (weight + activation FP4) + # Quantization mode: "w4a16" (weight-only NVFP4), "w4a4" (NVFP4 weight + activation), + # or "mxfp4_experts" (weight-only MXFP4 on routed MoE experts; keep TE fp8_autocast on) mode: "w4a16" # Group size for blockwise quantization (NVFP4 requires 16) group_size: 16 @@ -156,3 +157,5 @@ qat: activation_observer: null # Path to quantization config JSON for vLLM weight export quantization_config_path: null + # Let weight-only fake-quant run under TE fp8_autocast (real FP8 GEMM); lifts modelopt's guard. + bypass_te_fp8_assert: false diff --git a/verl/utils/modelopt/qat_utils.py b/verl/utils/modelopt/qat_utils.py index ca002c5132c..cdbc44b8d57 100644 --- a/verl/utils/modelopt/qat_utils.py +++ b/verl/utils/modelopt/qat_utils.py @@ -15,6 +15,56 @@ """High-level QAT workflow helpers for Megatron backend.""" +import logging +import os + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO")) # Ray workers keep root at WARNING + + +def _check_mxfp4_experts(modules): + """Fail fast if the experts-only MXFP4 config did not land where it should. + + (a) at least one routed-expert weight_quantizer enabled, (b) nothing else enabled (a mis-ordered + cfg leaves every other linear on modelopt's default INT8 fake-quant), (c) fake-quant is the + identity on the freshly loaded expert weights, which are already on the MXFP4 grid. + """ + import torch + from modelopt.torch.quantization.nn import TensorQuantizer + + n_expert_w, others, identity = 0, [], None + for m in modules: + for name, q in m.named_modules(): + if not isinstance(q, TensorQuantizer) or not getattr(q, "is_enabled", True): + continue + if name.endswith("weight_quantizer") and ".mlp.experts." in name: + n_expert_w += 1 + else: + others.append(name) + if identity is None: + for name, sub in m.named_modules(): + w0, wq = getattr(sub, "weight0", None), getattr(sub, "weight_quantizer", None) + if w0 is not None and wq is not None and ".mlp.experts." in name: + try: + with torch.no_grad(): + fq = wq(w0) + identity = (name, bool(torch.equal(fq, w0)), (fq.float() - w0.float()).abs().max().item()) + except Exception as e: # diagnostic only (e.g. weights not materialized yet) + identity = (name, "skipped", repr(e)[:120]) + break + logger.info( + "[QAT mxfp4_experts] enabled expert weight_quantizers=%d, other enabled quantizers=%d, " + "identity_check(name, equal, max_abs_diff)=%s", + n_expert_w, + len(others), + identity, + ) + has_experts = any(".mlp.experts." in n for m in modules for n, _ in m.named_modules()) + if has_experts and n_expert_w == 0: + raise RuntimeError("QAT mxfp4_experts: no routed-expert weight_quantizer enabled -- pattern did not match") + if others: + raise RuntimeError(f"QAT mxfp4_experts: {len(others)} unexpected quantizers enabled, e.g. {others[:3]}") + def _get_qat_field(qat_config, key, default=None): """Extract a field from qat_config, supporting both dict and object-style access.""" @@ -34,6 +84,8 @@ def apply_qat_to_modules(modules, qat_config): for i in range(len(modules)): modules[i] = apply_qat(modules[i], qat_mode, ignore_patterns=ignore_patterns) + if qat_mode == "mxfp4_experts": + _check_mxfp4_experts(modules) return modules diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py index dee7274099c..b16836f8209 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -21,6 +21,23 @@ import torch.nn as nn from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg +# MXFP4 = E2M1 weights with one E8M0 (power-of-two) scale per 32 contiguous K elements: the on-disk +# format of DeepSeek-V4-Flash routed experts and the format Megatron-Bridge re-quantizes them to on +# every vLLM weight sync, so training fake-quant and rollout weights share one grid. +_MXFP4_WEIGHT_CFG = { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, +} + +# DeepSeek-V4 report: "FP4 quantization-aware training for MoE expert weights" -- routed experts, +# weights only. Megatron paths are decoder.layers.N.mlp.experts.linear_fc{1,2}; the glob does not +# match mlp.shared_experts, the router, attention, or any input/output quantizer. The leading "*" +# disable is load-bearing: an unmatched TensorQuantizer defaults to ENABLED per-tensor INT8. +_MXFP4_EXPERTS_QUANT_CFG = [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*mlp.experts*weight_quantizer", "cfg": _MXFP4_WEIGHT_CFG, "enable": True}, +] + _NVFP4_W4A16_QUANTIZER_CFG = { "*weight_quantizer": { "num_bits": (2, 1), @@ -52,21 +69,28 @@ def build_quantize_config( ignore_patterns: list[str] | None = None, ) -> dict: """Build a complete ModelOpt quantization config for ``mtq.quantize``.""" - if qat_mode != "w4a16": - raise ValueError(f"Only 'w4a16' is supported, got: {qat_mode}") + if qat_mode == "w4a16": + quant_cfg = mtq.normalize_quant_cfg_list(_NVFP4_W4A16_QUANTIZER_CFG) + algorithm = "max" + elif qat_mode == "mxfp4_experts": + quant_cfg = copy.deepcopy(_MXFP4_EXPERTS_QUANT_CFG) + # Dynamic block quantization derives its scale per call; modelopt asserts it is never + # calibrated, so no algorithm / forward_loop. + algorithm = None + else: + raise ValueError(f"Only 'w4a16' and 'mxfp4_experts' are supported, got: {qat_mode}") if ignore_patterns is None: ignore_patterns = [] ignore_cfg = _ignore_patterns_to_quant_cfg(ignore_patterns) - quant_cfg = mtq.normalize_quant_cfg_list(_NVFP4_W4A16_QUANTIZER_CFG) disabled_cfg = copy.deepcopy(_default_disabled_quantizer_cfg) if isinstance(disabled_cfg, dict): disabled_cfg = mtq.normalize_quant_cfg_list(disabled_cfg) quant_cfg.extend(disabled_cfg) quant_cfg.extend(ignore_cfg) - return {"quant_cfg": quant_cfg, "algorithm": "max"} + return {"quant_cfg": quant_cfg, "algorithm": algorithm} def apply_qat( diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index 63783752090..d2e5e0a1150 100644 --- a/verl/workers/config/engine.py +++ b/verl/workers/config/engine.py @@ -142,6 +142,9 @@ class QATEngineConfig(BaseConfig): ignore_patterns: list[str] = field(default_factory=lambda: ["lm_head", "embed_tokens", "re:.*mlp.gate$"]) activation_observer: str = "static_minmax" quantization_config_path: Optional[str] = None + # Let weight-only fake-quant run under TE fp8_autocast (real FP8 GEMM). modelopt asserts against + # this combination by default; see MegatronEngine.initialize. + bypass_te_fp8_assert: bool = False @dataclass diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index d8e5e2e16d1..b2f287cb3fa 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -389,7 +389,11 @@ def _build_tf_config(self): else: provider_overrides["enable_routing_replay"] = True - if self._qat_enabled: + # Models with an experimental_attention_variant (DeepSeek-V4 "dsv4_hybrid", GLM5 "dsa", + # Qwen3-Next "gated_delta_net") install their own layer spec in provider_bridge(); the generic + # modelopt GPT spec would replace it and build a different model. modelopt swaps + # QuantModules by class, so the spec change is not required for quantization. + if self._qat_enabled and not getattr(provider, "experimental_attention_variant", None): from megatron.bridge.models.gpt_provider import modelopt_transformer_layer_spec provider.transformer_layer_spec = modelopt_transformer_layer_spec @@ -581,6 +585,19 @@ def initialize(self): self.module = self._build_megatron_module() if self._qat_enabled and not self.engine_config.forward_only: + if getattr(self._qat_config, "bypass_te_fp8_assert", False): + # Weight-only fake-quant feeding TE's real FP8 GEMM is the intended stack (DeepSeek-V4 does + # FP4 QAT inside FP8 training). modelopt refuses any quantizer under fp8_autocast; lift that + # guard at runtime. Activations are quantized once, by TE. modelopt's call sites resolve + # this module global by bare name at call time, so reassigning it is sufficient. + import modelopt.torch.quantization.plugins.transformer_engine as _mo_te + + if hasattr(_mo_te, "_assert_te_fp8_enabled"): + _mo_te._assert_te_fp8_enabled = lambda: None + logger.warning( + "QAT: bypassing modelopt _assert_te_fp8_enabled " + "(weight-only fake-quant + TE FP8 GEMM)" + ) from verl.utils.modelopt import apply_qat_to_modules self.module = apply_qat_to_modules(self.module, self._qat_config) @@ -1055,7 +1072,9 @@ def get_per_tensor_param(self, base_sync_done=False, **kwargs): ) # QAT: process weights through QATWeightExporter for quantized weight sync to vLLM - if self._qat_enabled: + # QATWeightExporter re-quantizes to NVFP4 for vLLM. MXFP4 experts are re-quantized by the bridge + # itself (_dsv4_use_mxfp4_export) on every sync, so that path is skipped for mxfp4_experts. + if self._qat_enabled and self._qat_config.mode in ("w4a16", "w4a4"): from verl.utils.modelopt import export_qat_weights per_tensor_param = export_qat_weights(per_tensor_param, self.module, self._qat_config.mode, self.bridge) From ac61712baab0cd3298d44cde50420ec5845abc6b Mon Sep 17 00:00:00 2001 From: sopyang Date: Thu, 10 Sep 2026 01:38:13 -0700 Subject: [PATCH 2/6] [example] add DeepSeek-V4-Flash GRPO recipe with MXFP4 expert QAT Mirrors run_deepseek_v4_flash_megatron.sh with TE fp8 blockwise training, fp8 vLLM rollout and the optional `mxfp4_experts` QAT (off by default). Documents why the FP4 step is simulated on Hopper while the FP8 GEMM is not, and why filter_overlong_prompts must stay off for DeepSeek-V4. Co-Authored-By: Claude Opus 5 (1M context) --- .../run_deepseek_v4_flash_mcore_fp4qat.sh | 313 ++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100755 examples/grpo_trainer/run_deepseek_v4_flash_mcore_fp4qat.sh diff --git a/examples/grpo_trainer/run_deepseek_v4_flash_mcore_fp4qat.sh b/examples/grpo_trainer/run_deepseek_v4_flash_mcore_fp4qat.sh new file mode 100755 index 00000000000..84ab1614fc3 --- /dev/null +++ b/examples/grpo_trainer/run_deepseek_v4_flash_mcore_fp4qat.sh @@ -0,0 +1,313 @@ +#!/usr/bin/env bash +# GRPO | DeepSeek-V4-Flash | vLLM fp8 rollout | Megatron fp8 training + MXFP4 expert QAT | NVIDIA GPUs +# +# What this adds over run_deepseek_v4_flash_megatron.sh +# ----------------------------------------------------- +# The DeepSeek-V4 technical report states that, during post-training, "we incorporate FP4 +# quantization-aware training for MoE expert weights". This script reproduces that: a weight-only +# fake-quant onto the MXFP4 grid (E2M1 values, one E8M0 power-of-two scale per 32 contiguous K +# elements) is applied to the ROUTED experts only, while the GEMMs keep running in real FP8 through +# TransformerEngine's blockwise autocast. +# +# weight -> MXFP4 fake-quant (bf16 container on the FP4 grid) -> TE cast to FP8 E4M3 128x128 +# input -> TE cast to FP8 E4M3 1x128 +# -> real FP8 tensor-core GEMM +# +# The FP4 step is simulated because Hopper has no FP4 tensor cores; the FP8 GEMM is not simulated. +# On Blackwell the same configuration is the natural starting point for a native FP4 kernel. +# +# Why this matches the checkpoint. DeepSeek-V4-Flash ships its routed experts already quantized to +# MXFP4 (packed E2M1 + F8_E8M0 scales, block 32) while every other quantized tensor is FP8 E4M3 with +# 128x128 scales. Megatron-Bridge re-quantizes exactly those expert tensors back to MXFP4 on every +# weight sync to vLLM, so the training forward and the rollout weights share one grid. +# +# Knobs +# QAT_ENABLE=True turn the fake-quant on (default False, i.e. plain fp8 training) +# QAT_MODE=mxfp4_experts verl/utils/modelopt/quantize.py; scopes to *mlp.experts* +# QAT_BYPASS_TE_FP8_ASSERT=True ModelOpt refuses to run any quantizer under TE's fp8_autocast. +# That guard exists to keep pure simulation faithful; here the +# stacking is intentional (FP4 weights feeding a real FP8 GEMM), so +# it is lifted at runtime. Activations are quantized once, by TE -- +# do not also enable ModelOpt's input_quantizer. +# +# Requires nvidia-modelopt (its MX fake-quant CUDA extension is JIT-built on first use, so nvcc must +# be available on the workers) and Megatron-Bridge importable on every node. +# +# With: +# - Megatron-Bridge: https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/c7774d44d4b3101dc6bdf8c8d38a32e909e1ea11 +# - Megatron-LM: https://github.com/NVIDIA/Megatron-LM/commit/1ff25ca7e339fe521165da7f4373d9f52e7af436 + +set -euo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NCCL_NVLS_ENABLE=0 +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export VLLM_USE_V1=1 +export PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +# Keep TE and vLLM on FP32 block scales rather than ue8m0, so both sides agree on the FP8 grid. +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 +export VLLM_USE_DEEP_GEMM_E8M0=0 +# The ModelOpt MX kernel is JIT-compiled for the visible device; pin it so a future image listing +# many architectures does not multiply the one-off nvcc time. +export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-9.0} + +############################### configs ################################ + +MODEL_PATH=${MODEL_PATH:-$HDFS_ROOT/model/DeepSeek-V4-Flash} +NNODES=${NNODES:-16} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} + +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +DATALOADER_NUM_WORKERS=${DATALOADER_NUM_WORKERS:-8} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH))} + +USE_KL_IN_REWARD=${USE_KL_IN_REWARD:-False} +KL_COEF=${KL_COEF:-0.001} + +ACTOR_LR=${ACTOR_LR:-1e-6} +OPTIMIZER_OFFLOAD_FRACTION=${OPTIMIZER_OFFLOAD_FRACTION:-1.0} + +ACTOR_TP=${ACTOR_TP:-1} +ACTOR_PP=${ACTOR_PP:-8} +ACTOR_VPP=${ACTOR_VPP:-null} +ACTOR_EP=${ACTOR_EP:-16} +ACTOR_ETP=${ACTOR_ETP:-1} +ACTOR_CP=${ACTOR_CP:-1} +PIPELINE_MODEL_PARALLEL_LAYOUT=${PIPELINE_MODEL_PARALLEL_LAYOUT:-"Et*6|t*6|t*6|t*5|t*5|t*5|t*5|t*5L"} + +REF_TP=${REF_TP:-${ACTOR_TP}} +REF_PP=${REF_PP:-${ACTOR_PP}} +REF_VPP=${REF_VPP:-${ACTOR_VPP}} +REF_EP=${REF_EP:-${ACTOR_EP}} +REF_ETP=${REF_ETP:-${ACTOR_ETP}} +REF_CP=${REF_CP:-${ACTOR_CP}} + +ROLLOUT_EP=${ROLLOUT_EP:-8} +ROLLOUT_N=${ROLLOUT_N:-8} +ROLLOUT_ENFORCE_EAGER=${ROLLOUT_ENFORCE_EAGER:-False} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.40} +ROLLOUT_MAX_MODEL_LEN=${ROLLOUT_MAX_MODEL_LEN:-${PPO_MAX_TOKEN_LEN_PER_GPU}} +ROLLOUT_MAX_NUM_BATCHED_TOKENS=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-${PPO_MAX_TOKEN_LEN_PER_GPU}} +ROLLOUT_KV_CACHE_DTYPE=${ROLLOUT_KV_CACHE_DTYPE:-fp8} +ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB:-512} +ROUTER_REPLAY_MODE=${ROUTER_REPLAY_MODE:-R3} + +# Truncated importance sampling over the rollout/trainer logprob ratio. Set to "off" to disable. +ROLLOUT_IS_MODE=${ROLLOUT_IS_MODE:-token} +ROLLOUT_IS_THRESHOLD=${ROLLOUT_IS_THRESHOLD:-2.0} + +# MXFP4 quantization-aware training on the routed MoE experts. See the header. +QAT_ENABLE=${QAT_ENABLE:-False} +QAT_MODE=${QAT_MODE:-mxfp4_experts} +QAT_BYPASS_TE_FP8_ASSERT=${QAT_BYPASS_TE_FP8_ASSERT:-True} + +ALL_OFFLOAD=${ALL_OFFLOAD:-True} +TOTAL_EPOCHS=${TOTAL_EPOCHS:-1} +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-30} +SAVE_FREQ=${SAVE_FREQ:--1} +TEST_FREQ=${TEST_FREQ:--1} +# Directory for raw rollout dumps (prompt/response/score as jsonl); "null" disables them. +ROLLOUT_DATA_DIR=${ROLLOUT_DATA_DIR:-null} + +PROJECT_NAME=${PROJECT_NAME:-verl_dsv4_flash} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-deepseek_v4_flash_grpo_fp8_mxfp4qat} +CKPTS_DIR=${CKPTS_DIR:-"${HOME}/verl/ckpts/${PROJECT_NAME}/${EXPERIMENT_NAME}"} + +TRAIN_FILE=${TRAIN_FILE:-$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k/data/dapo-math-17k.parquet} +TEST_FILE=${TEST_FILE:-$DATA_ROOT/dataset/aime25_test.parquet} +OVERLONG_BUFFER_LEN=${OVERLONG_BUFFER_LEN:-${MAX_RESPONSE_LENGTH}} +OVERLONG_BUFFER_ENABLE=${OVERLONG_BUFFER_ENABLE:-False} +OVERLONG_PENALTY_FACTOR=${OVERLONG_PENALTY_FACTOR:-1.0} + +########################### parameter arrays ########################### + +ENABLE_THINKING=${ENABLE_THINKING:-True} + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=${USE_KL_IN_REWARD} + algorithm.kl_ctrl.kl_coef=${KL_COEF} +) + +DATA=( + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.prompt_key=prompt + data.return_raw_chat=True + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + # DeepSeek-V4 renders prompts through verl's own continuous-token builder rather than a jinja + # chat template, and the official checkpoints ship none. Length filtering is the one remaining + # path that would still call apply_chat_template, so keep it off. + data.filter_overlong_prompts=False + data.truncation=error + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS} + +data.apply_chat_template_kwargs.enable_thinking=${ENABLE_THINKING} +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} + actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} + actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.use_distributed_optimizer=True + actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTER_REPLAY_MODE} + ++actor_rollout_ref.actor.megatron.override_transformer_config.apply_dsa_kernel_fusion=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_use_sparse_loss=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_loss_coeff=0.0 + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + ++actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_mhc=False + # TE fp8 GEMM autocast: the "real FP8" half of the recipe. It stays on together with QAT. + +actor_rollout_ref.actor.megatron.override_transformer_config.fp8="hybrid" + +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe="blockwise" + +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe="blockwise" + "++actor_rollout_ref.actor.megatron.override_transformer_config.pipeline_model_parallel_layout='${PIPELINE_MODEL_PARALLEL_LAYOUT}'" +) + +# Context parallelism needs three extra transformer-config settings for DeepSeek-V4 on top of +# `context_parallel_size`. Each of them is enforced by Megatron-Core, so without them the run +# aborts at model build or in the first attention forward. They are appended only when CP > 1. +# +# cp_partition_mode=contiguous +# DSv4 attention requires every CP rank to own ONE consecutive interval of the packed THD +# buffer. Megatron-Core defaults to "zigzag" and raises +# "DSv4 Hybrid with CP requires cp_partition_mode='contiguous'." +# sequence_packing_scheduler=dp_balanced +# Megatron-Core asserts "DSv4 Hybrid with CP requires a sequence_packing_scheduler for THD +# inputs." Needs Transformer Engine >= 2.9. +# max_seqlen_per_dp_cp_rank +# Documented as max sequence length / cp_size; it drives how sub-samples are assigned to +# each DPxCP rank. +CP_ARGS=() +if [ "${ACTOR_CP}" -gt 1 ]; then + CP_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.cp_partition_mode=contiguous + ++actor_rollout_ref.actor.megatron.override_transformer_config.sequence_packing_scheduler=dp_balanced + ++actor_rollout_ref.actor.megatron.override_transformer_config.max_seqlen_per_dp_cp_rank=$(((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH) / ACTOR_CP)) + ) +fi + +QAT_ARGS=() +if [ "${QAT_ENABLE}" = "True" ]; then + QAT_ARGS=( + ++actor_rollout_ref.actor.megatron.qat.enable=True + ++actor_rollout_ref.actor.megatron.qat.mode=${QAT_MODE} + ++actor_rollout_ref.actor.megatron.qat.bypass_te_fp8_assert=${QAT_BYPASS_TE_FP8_ASSERT} + # rollout.qat is an oc.select alias of actor.megatron.qat in the generated config. The vLLM + # worker must not build quantizers: it consumes the MXFP4 weights the bridge exports. + ++actor_rollout_ref.rollout.qat.enable=False + ) +fi + +TIS_ARGS=() +if [ "${ROLLOUT_IS_MODE}" != "off" ]; then + TIS_ARGS=( + algorithm.rollout_correction.rollout_is=${ROLLOUT_IS_MODE} + algorithm.rollout_correction.rollout_is_threshold=${ROLLOUT_IS_THRESHOLD} + ) +fi + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=1 + actor_rollout_ref.rollout.data_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.expert_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.enforce_eager=${ROLLOUT_ENFORCE_EAGER} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.enable_rollout_routing_replay=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} + actor_rollout_ref.rollout.max_model_len=${ROLLOUT_MAX_MODEL_LEN} + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB} + +actor_rollout_ref.rollout.engine_kwargs.vllm.kv_cache_dtype=${ROLLOUT_KV_CACHE_DTYPE} + +actor_rollout_ref.rollout.quantization=fp8 +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=${OVERLONG_BUFFER_ENABLE} + +reward.reward_kwargs.overlong_buffer_cfg.len=${OVERLONG_BUFFER_LEN} + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${OVERLONG_PENALTY_FACTOR} + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${MAX_RESPONSE_LENGTH} +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.total_epochs=${TOTAL_EPOCHS} + trainer.total_training_steps=${TOTAL_TRAINING_STEPS} + trainer.resume_mode=auto + trainer.val_before_train=False + trainer.log_val_generations=0 + trainer.rollout_data_dir=${ROLLOUT_DATA_DIR} + trainer.default_local_dir="${CKPTS_DIR}" +) + +EXTRA=( + actor_rollout_ref.nccl_timeout=3600 + model_engine=megatron +) + +########################### launch ########################### + +# uv (set VERL_USE_UV=0 for system python): the GPU vllm x megatron driver and every Ray worker +# (runtime_env.py_executable) run through `uv run` on the matching extras of the committed uv.lock. +# Run from the verl repo root. +LAUNCH=(python3) +RAY=(ray_kwargs.ray_init.runtime_env.py_executable=null) +if [ "${VERL_USE_UV:-1}" != 0 ] && [ "${DEVICE:-gpu}" = gpu ]; then + LAUNCH=(uv run --frozen --all-packages --extra vllm --extra megatron python3) + RAY=(ray_kwargs.ray_init.runtime_env.py_executable="uv -v run --frozen --all-packages --extra vllm --extra megatron") +fi +"${LAUNCH[@]}" -m verl.trainer.main_ppo \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${CP_ARGS[@]}" \ + "${QAT_ARGS[@]}" \ + "${TIS_ARGS[@]}" \ + "${ROLLOUT[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "${RAY[@]}" \ + "$@" From 2628219dd914f9a730e71237b8b815394f0e6964 Mon Sep 17 00:00:00 2001 From: Sophia Yang Date: Mon, 14 Sep 2026 08:43:22 -0700 Subject: [PATCH 3/6] [megatron] fix MXFP4 QAT checkpoints and add length experiments Preserve native Megatron checkpoint methods around stateless expert MXFP4 quantization, and restore offloaded optimizer shards before loading HDO state. Add independent length-only and MXFP4 expert QAT recipes with a 4096-token response budget, 40 total steps and validation every 10 steps. Require model and dataset locations from the environment and use portable output defaults. Validation: checkpoint regressions passed (6 and 7 tests); the original length-only and QAT runs each completed 40 steps including save/resume. Published launch arguments match both completed experiments at 20/40 steps. Shell syntax, required-input checks and credential/path scans passed. Two pre-existing repository checks are excluded: stale generated reference configuration and naming of an older example. New example names pass. Co-authored-by: OpenAI Codex --- ..._deepseek_v4_flash_length_4096_megatron.sh | 287 +++++++++++++++++ ...pseek_v4_flash_length_4096_qat_megatron.sh | 290 ++++++++++++++++++ .../megatron/test_hdo_checkpoint_resume.py | 159 ++++++++++ tests/utils/modelopt/test_mxfp4_checkpoint.py | 229 ++++++++++++++ verl/utils/modelopt/checkpoint.py | 56 ++++ verl/utils/modelopt/quantize.py | 8 +- .../engine/megatron/transformer_impl.py | 9 +- 7 files changed, 1035 insertions(+), 3 deletions(-) create mode 100755 examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh create mode 100755 examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh create mode 100644 tests/utils/megatron/test_hdo_checkpoint_resume.py create mode 100644 tests/utils/modelopt/test_mxfp4_checkpoint.py create mode 100644 verl/utils/modelopt/checkpoint.py diff --git a/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh new file mode 100755 index 00000000000..902a42b4316 --- /dev/null +++ b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh @@ -0,0 +1,287 @@ +#!/usr/bin/env bash +# DeepSeek-V4-Flash-0731: response-length experiment, 40 steps, validation every 10. +# Standalone length-only baseline; the MXFP4 QAT arm is a separate script. +# Plan A: prompt/response 1024/4096, overlong buffer 512 with penalty 1.0, +# clip high 0.28 / dual clip 10.0, no thinking, no QAT; retain TE FP8 training. +# Set MODEL_PATH, TRAIN_FILE and TEST_FILE in the launch environment. +# Optional: set OUTPUT_DIR or override each checkpoint/rollout/validation directory. +# For two segments, use TOTAL_TRAINING_STEPS=20, then 40 with the same checkpoint directory. +# Requires DeepSeek-V4 support in Megatron-Bridge, Megatron-Core and vLLM. + +set -euo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NCCL_NVLS_ENABLE=0 +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export VLLM_USE_V1=1 +export PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 # Use FP32 scales for TE block-wise FP8 +export VLLM_USE_DEEP_GEMM_E8M0=0 # Use FP32 scales for vLLM with DeepGEMM + +############################### configs ################################ + +: "${MODEL_PATH:?Set MODEL_PATH to the model checkpoint directory}" +: "${TRAIN_FILE:?Set TRAIN_FILE to the training parquet file}" +: "${TEST_FILE:?Set TEST_FILE to the validation parquet file}" +OUTPUT_DIR=${OUTPUT_DIR:-./outputs} +NNODES=${NNODES:-16} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +DATALOADER_NUM_WORKERS=${DATALOADER_NUM_WORKERS:-8} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-4096} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-5120} + +USE_KL_IN_REWARD=${USE_KL_IN_REWARD:-False} +KL_COEF=${KL_COEF:-0.001} +CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-0.2} +CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} +CLIP_RATIO_C=${CLIP_RATIO_C:-10.0} +ACTOR_LR=${ACTOR_LR:-1e-6} +OPTIMIZER_OFFLOAD_FRACTION=${OPTIMIZER_OFFLOAD_FRACTION:-1.0} + +ACTOR_TP=${ACTOR_TP:-1} +ACTOR_PP=${ACTOR_PP:-8} +ACTOR_VPP=${ACTOR_VPP:-null} +ACTOR_EP=${ACTOR_EP:-16} +ACTOR_ETP=${ACTOR_ETP:-1} +ACTOR_CP=${ACTOR_CP:-1} +PIPELINE_MODEL_PARALLEL_LAYOUT=${PIPELINE_MODEL_PARALLEL_LAYOUT:-"Et*6|t*6|t*6|t*5|t*5|t*5|t*5|t*5L"} +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +ROLLOUT_EP=${ROLLOUT_EP:-8} +ROLLOUT_N=${ROLLOUT_N:-8} +ROLLOUT_N_VAL=${ROLLOUT_N_VAL:-16} +ROLLOUT_ENFORCE_EAGER=${ROLLOUT_ENFORCE_EAGER:-False} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.4} +ROLLOUT_MAX_MODEL_LEN=${ROLLOUT_MAX_MODEL_LEN:-5120} +# Keep the proven prefill budget independent of the longer response/context cap. +ROLLOUT_MAX_NUM_BATCHED_TOKENS=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-3072} +ROLLOUT_KV_CACHE_DTYPE=${ROLLOUT_KV_CACHE_DTYPE:-fp8} +ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB:-512} +ROUTER_REPLAY_MODE=${ROUTER_REPLAY_MODE:-R3} +ROLLOUT_IS_MODE=${ROLLOUT_IS_MODE:-token} +ROLLOUT_IS_THRESHOLD=${ROLLOUT_IS_THRESHOLD:-2.0} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} +# The submitter sets this to 20 for segment 1 and 40 for the resumed segment. +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-40} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-10} +PROJECT_NAME=${PROJECT_NAME:-verl_deepseek_v4_flash} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-dsv4-0731-noqat-len4096-40step-val10} +CKPTS_DIR=${CKPTS_DIR:-$OUTPUT_DIR/checkpoints/$PROJECT_NAME/$EXPERIMENT_NAME} +ROLLOUT_DATA_DIR=${ROLLOUT_DATA_DIR:-$OUTPUT_DIR/rollout_dumps/$EXPERIMENT_NAME} +VALIDATION_DATA_DIR=${VALIDATION_DATA_DIR:-$OUTPUT_DIR/validation_dumps/$EXPERIMENT_NAME} +OVERLONG_BUFFER_LEN=${OVERLONG_BUFFER_LEN:-512} +OVERLONG_BUFFER_ENABLE=${OVERLONG_BUFFER_ENABLE:-True} +OVERLONG_PENALTY_FACTOR=${OVERLONG_PENALTY_FACTOR:-1.0} +ENABLE_THINKING=${ENABLE_THINKING:-False} + +########################### parameter arrays ########################### + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=${USE_KL_IN_REWARD} + algorithm.kl_ctrl.kl_coef=${KL_COEF} +) + +DATA=( + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.prompt_key=prompt + data.return_raw_chat=True + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=False + data.filter_overlong_prompts_workers=64 + data.truncation=error + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS} + +data.apply_chat_template_kwargs.enable_thinking=${ENABLE_THINKING} +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + # actor_rollout_ref.actor.policy_loss.loss_mode=${LOSS_MODE} + # actor_rollout_ref.actor.loss_agg_mode=${LOSS_AGG_MODE} + # actor_rollout_ref.actor.use_kl_loss=${USE_KL_LOSS} + # actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} + actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} + actor_rollout_ref.actor.clip_ratio_c=${CLIP_RATIO_C} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} + actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} + actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} + # grad_offload removed upstream by verl 559c337a (#7544, BREAKING) as unused -- + # setting it now raises ConfigAttributeError: Key 'grad_offload' is not in struct. + actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTER_REPLAY_MODE} + ++actor_rollout_ref.actor.megatron.override_transformer_config.apply_dsa_kernel_fusion=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_use_sparse_loss=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_loss_coeff=0.0 + ++actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_mhc=False + # moe_router_fusion: TE fused-router backward conflicts with router_replay R3 (inplace-modified + # routing map in autograd) — RuntimeError at te/pytorch/router.py:72. Keep OFF with R3. + # ++actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_fusion=True + # ++actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_padding_for_quantization=True + +actor_rollout_ref.actor.megatron.override_transformer_config.fp8="hybrid" + +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe="blockwise" + +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe="blockwise" + # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_param=True + # +actor_rollout_ref.actor.megatron.override_ddp_config.fp8_param_gather=True + # +actor_rollout_ref.actor.megatron.override_ddp_config.overlap_grad_reduce=True + # +actor_rollout_ref.actor.megatron.override_ddp_config.overlap_param_gather=True + actor_rollout_ref.actor.megatron.use_distributed_optimizer=True + "++actor_rollout_ref.actor.megatron.override_transformer_config.pipeline_model_parallel_layout='${PIPELINE_MODEL_PARALLEL_LAYOUT}'" +) + +if [ "${RECOMPUTE_GRANULARITY}" = "selective" ]; then + RECOMPUTE_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=selective + "++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_modules=[mla_up_proj]" + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=null + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=null + ) +else + RECOMPUTE_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + ) +fi + +# Context parallelism needs three extra transformer-config settings for DeepSeek-V4 on top of +# `context_parallel_size`. Each of them is enforced by Megatron-Core, so without them the run +# aborts at model build or in the first attention forward. They are appended only when CP > 1. +# +# cp_partition_mode=contiguous +# DSv4 attention requires every CP rank to own ONE consecutive interval of the packed THD +# buffer. Megatron-Core defaults to "zigzag" and raises +# "DSv4 Hybrid with CP requires cp_partition_mode='contiguous'." +# sequence_packing_scheduler=dp_balanced +# Megatron-Core asserts "DSv4 Hybrid with CP requires a sequence_packing_scheduler for THD +# inputs." Needs Transformer Engine >= 2.9. +# max_seqlen_per_dp_cp_rank +# Documented as max sequence length / cp_size; it drives how sub-samples are assigned to +# each DPxCP rank. +CP_ARGS=() +if [ "${ACTOR_CP}" -gt 1 ]; then + CP_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.cp_partition_mode=contiguous + ++actor_rollout_ref.actor.megatron.override_transformer_config.sequence_packing_scheduler=dp_balanced + ++actor_rollout_ref.actor.megatron.override_transformer_config.max_seqlen_per_dp_cp_rank=$(((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH) / ACTOR_CP)) + ) +fi + +# 截断重要性采样, 对 rollout/trainer 的 logprob 比值做裁剪。 +TIS_ARGS=() +if [ "${ROLLOUT_IS_MODE}" != "off" ]; then + TIS_ARGS=( + algorithm.rollout_correction.rollout_is=${ROLLOUT_IS_MODE} + algorithm.rollout_correction.rollout_is_threshold=${ROLLOUT_IS_THRESHOLD} + ) +fi + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=1 + actor_rollout_ref.rollout.data_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.expert_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.enforce_eager=${ROLLOUT_ENFORCE_EAGER} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.n=${ROLLOUT_N_VAL} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.enable_rollout_routing_replay=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} + actor_rollout_ref.rollout.max_model_len=${ROLLOUT_MAX_MODEL_LEN} + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB} + +actor_rollout_ref.rollout.engine_kwargs.vllm.kv_cache_dtype=${ROLLOUT_KV_CACHE_DTYPE} + +actor_rollout_ref.rollout.quantization=fp8 +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=${OVERLONG_BUFFER_ENABLE} + +reward.reward_kwargs.overlong_buffer_cfg.len=${OVERLONG_BUFFER_LEN} + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${OVERLONG_PENALTY_FACTOR} + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${MAX_RESPONSE_LENGTH} +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.max_actor_ckpt_to_keep=2 + trainer.total_epochs=${TOTAL_EPOCHS} + trainer.total_training_steps=${TOTAL_TRAINING_STEPS} + trainer.resume_mode=auto + trainer.val_before_train=False + trainer.log_val_generations=100 + trainer.rollout_data_dir="${ROLLOUT_DATA_DIR}" + trainer.validation_data_dir="${VALIDATION_DATA_DIR}" + trainer.default_local_dir="${CKPTS_DIR}" +) + +EXTRA=( + actor_rollout_ref.nccl_timeout=3600 + model_engine=megatron +) + +# Explicitly disable QAT for both actor and rollout in this length-only recipe. +QAT_ARGS=( + ++actor_rollout_ref.actor.megatron.qat.enable=False + ++actor_rollout_ref.rollout.qat.enable=False +) + +########################### launch ########################### + +python3 -m verl.trainer.main_ppo \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${RECOMPUTE_ARGS[@]}" \ + "${CP_ARGS[@]}" \ + "${TIS_ARGS[@]}" \ + "${ROLLOUT[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "${QAT_ARGS[@]}" \ + "$@" diff --git a/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh new file mode 100755 index 00000000000..2ec004d952e --- /dev/null +++ b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh @@ -0,0 +1,290 @@ +#!/usr/bin/env bash +# DeepSeek-V4-Flash-0731: MXFP4 QAT, response 4096, 40 steps, validation every 10. +# Standalone MXFP4 QAT arm; the length-only baseline is a separate script. +# Plan A: prompt/response 1024/4096, overlong buffer 512 with penalty 1.0, +# clip high 0.28 / dual clip 10.0, no thinking; expert MXFP4 QAT plus TE FP8 training. +# Set MODEL_PATH, TRAIN_FILE and TEST_FILE in the launch environment. +# Optional: set OUTPUT_DIR or override each checkpoint/rollout/validation directory. +# For two segments, use TOTAL_TRAINING_STEPS=20, then 40 with the same checkpoint directory. +# Requires DeepSeek-V4 support in Megatron-Bridge, Megatron-Core and vLLM. + +set -euo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NCCL_NVLS_ENABLE=0 +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export VLLM_USE_V1=1 +export PYTHONUNBUFFERED=1 +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 # Use FP32 scales for TE block-wise FP8 +export VLLM_USE_DEEP_GEMM_E8M0=0 # Use FP32 scales for vLLM with DeepGEMM + +############################### configs ################################ + +: "${MODEL_PATH:?Set MODEL_PATH to the model checkpoint directory}" +: "${TRAIN_FILE:?Set TRAIN_FILE to the training parquet file}" +: "${TEST_FILE:?Set TEST_FILE to the validation parquet file}" +OUTPUT_DIR=${OUTPUT_DIR:-./outputs} +NNODES=${NNODES:-16} +NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} +PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} +ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU:-1} +DATALOADER_NUM_WORKERS=${DATALOADER_NUM_WORKERS:-8} +MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-4096} +PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-5120} + +USE_KL_IN_REWARD=${USE_KL_IN_REWARD:-False} +KL_COEF=${KL_COEF:-0.001} +CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-0.2} +CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} +CLIP_RATIO_C=${CLIP_RATIO_C:-10.0} +ACTOR_LR=${ACTOR_LR:-1e-6} +OPTIMIZER_OFFLOAD_FRACTION=${OPTIMIZER_OFFLOAD_FRACTION:-1.0} + +ACTOR_TP=${ACTOR_TP:-1} +ACTOR_PP=${ACTOR_PP:-8} +ACTOR_VPP=${ACTOR_VPP:-null} +ACTOR_EP=${ACTOR_EP:-16} +ACTOR_ETP=${ACTOR_ETP:-1} +ACTOR_CP=${ACTOR_CP:-1} +PIPELINE_MODEL_PARALLEL_LAYOUT=${PIPELINE_MODEL_PARALLEL_LAYOUT:-"Et*6|t*6|t*6|t*5|t*5|t*5|t*5|t*5L"} +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} +ALL_OFFLOAD=${ALL_OFFLOAD:-True} + +ROLLOUT_EP=${ROLLOUT_EP:-8} +ROLLOUT_N=${ROLLOUT_N:-8} +ROLLOUT_N_VAL=${ROLLOUT_N_VAL:-16} +ROLLOUT_ENFORCE_EAGER=${ROLLOUT_ENFORCE_EAGER:-False} +ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.4} +ROLLOUT_MAX_MODEL_LEN=${ROLLOUT_MAX_MODEL_LEN:-5120} +# Keep the proven prefill budget independent of the longer response/context cap. +ROLLOUT_MAX_NUM_BATCHED_TOKENS=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-3072} +ROLLOUT_KV_CACHE_DTYPE=${ROLLOUT_KV_CACHE_DTYPE:-fp8} +ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB:-512} +ROUTER_REPLAY_MODE=${ROUTER_REPLAY_MODE:-R3} +ROLLOUT_IS_MODE=${ROLLOUT_IS_MODE:-token} +ROLLOUT_IS_THRESHOLD=${ROLLOUT_IS_THRESHOLD:-2.0} + +TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} +# The submitter sets this to 20 for segment 1 and 40 for the resumed segment. +TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-40} +SAVE_FREQ=${SAVE_FREQ:-20} +TEST_FREQ=${TEST_FREQ:-10} +PROJECT_NAME=${PROJECT_NAME:-verl_deepseek_v4_flash} +EXPERIMENT_NAME=${EXPERIMENT_NAME:-dsv4-0731-mxfp4qat-len4096-40step-val10} +CKPTS_DIR=${CKPTS_DIR:-$OUTPUT_DIR/checkpoints/$PROJECT_NAME/$EXPERIMENT_NAME} +ROLLOUT_DATA_DIR=${ROLLOUT_DATA_DIR:-$OUTPUT_DIR/rollout_dumps/$EXPERIMENT_NAME} +VALIDATION_DATA_DIR=${VALIDATION_DATA_DIR:-$OUTPUT_DIR/validation_dumps/$EXPERIMENT_NAME} +OVERLONG_BUFFER_LEN=${OVERLONG_BUFFER_LEN:-512} +OVERLONG_BUFFER_ENABLE=${OVERLONG_BUFFER_ENABLE:-True} +OVERLONG_PENALTY_FACTOR=${OVERLONG_PENALTY_FACTOR:-1.0} +ENABLE_THINKING=${ENABLE_THINKING:-False} + +########################### parameter arrays ########################### + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=${USE_KL_IN_REWARD} + algorithm.kl_ctrl.kl_coef=${KL_COEF} +) + +DATA=( + data.train_files="$TRAIN_FILE" + data.val_files="$TEST_FILE" + data.train_batch_size=${TRAIN_BATCH_SIZE} + data.prompt_key=prompt + data.return_raw_chat=True + data.max_prompt_length=${MAX_PROMPT_LENGTH} + data.max_response_length=${MAX_RESPONSE_LENGTH} + data.filter_overlong_prompts=False + data.filter_overlong_prompts_workers=64 + data.truncation=error + data.dataloader_num_workers=${DATALOADER_NUM_WORKERS} + +data.apply_chat_template_kwargs.enable_thinking=${ENABLE_THINKING} +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.use_fused_kernels=True + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.enable_gradient_checkpointing=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${ACTOR_LR} + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} + +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True + +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True + # actor_rollout_ref.actor.policy_loss.loss_mode=${LOSS_MODE} + # actor_rollout_ref.actor.loss_agg_mode=${LOSS_AGG_MODE} + # actor_rollout_ref.actor.use_kl_loss=${USE_KL_LOSS} + # actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} + actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} + actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} + actor_rollout_ref.actor.clip_ratio_c=${CLIP_RATIO_C} + actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} + actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} + actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} + actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} + actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} + # grad_offload removed upstream by verl 559c337a (#7544, BREAKING) as unused -- + # setting it now raises ConfigAttributeError: Key 'grad_offload' is not in struct. + actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTER_REPLAY_MODE} + ++actor_rollout_ref.actor.megatron.override_transformer_config.apply_dsa_kernel_fusion=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_use_sparse_loss=True + ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_loss_coeff=0.0 + ++actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_mhc=False + # moe_router_fusion: TE fused-router backward conflicts with router_replay R3 (inplace-modified + # routing map in autograd) — RuntimeError at te/pytorch/router.py:72. Keep OFF with R3. + # ++actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_fusion=True + # ++actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_padding_for_quantization=True + +actor_rollout_ref.actor.megatron.override_transformer_config.fp8="hybrid" + +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe="blockwise" + +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe="blockwise" + # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_param=True + # +actor_rollout_ref.actor.megatron.override_ddp_config.fp8_param_gather=True + # +actor_rollout_ref.actor.megatron.override_ddp_config.overlap_grad_reduce=True + # +actor_rollout_ref.actor.megatron.override_ddp_config.overlap_param_gather=True + actor_rollout_ref.actor.megatron.use_distributed_optimizer=True + "++actor_rollout_ref.actor.megatron.override_transformer_config.pipeline_model_parallel_layout='${PIPELINE_MODEL_PARALLEL_LAYOUT}'" +) + +if [ "${RECOMPUTE_GRANULARITY}" = "selective" ]; then + RECOMPUTE_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=selective + "++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_modules=[mla_up_proj]" + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=null + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=null + ) +else + RECOMPUTE_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + ) +fi + +# Context parallelism needs three extra transformer-config settings for DeepSeek-V4 on top of +# `context_parallel_size`. Each of them is enforced by Megatron-Core, so without them the run +# aborts at model build or in the first attention forward. They are appended only when CP > 1. +# +# cp_partition_mode=contiguous +# DSv4 attention requires every CP rank to own ONE consecutive interval of the packed THD +# buffer. Megatron-Core defaults to "zigzag" and raises +# "DSv4 Hybrid with CP requires cp_partition_mode='contiguous'." +# sequence_packing_scheduler=dp_balanced +# Megatron-Core asserts "DSv4 Hybrid with CP requires a sequence_packing_scheduler for THD +# inputs." Needs Transformer Engine >= 2.9. +# max_seqlen_per_dp_cp_rank +# Documented as max sequence length / cp_size; it drives how sub-samples are assigned to +# each DPxCP rank. +CP_ARGS=() +if [ "${ACTOR_CP}" -gt 1 ]; then + CP_ARGS=( + ++actor_rollout_ref.actor.megatron.override_transformer_config.cp_partition_mode=contiguous + ++actor_rollout_ref.actor.megatron.override_transformer_config.sequence_packing_scheduler=dp_balanced + ++actor_rollout_ref.actor.megatron.override_transformer_config.max_seqlen_per_dp_cp_rank=$(((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH) / ACTOR_CP)) + ) +fi + +# 截断重要性采样, 对 rollout/trainer 的 logprob 比值做裁剪。 +TIS_ARGS=() +if [ "${ROLLOUT_IS_MODE}" != "off" ]; then + TIS_ARGS=( + algorithm.rollout_correction.rollout_is=${ROLLOUT_IS_MODE} + algorithm.rollout_correction.rollout_is_threshold=${ROLLOUT_IS_THRESHOLD} + ) +fi + +ROLLOUT=( + actor_rollout_ref.rollout.name=vllm + actor_rollout_ref.rollout.tensor_model_parallel_size=1 + actor_rollout_ref.rollout.data_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.expert_parallel_size=${ROLLOUT_EP} + actor_rollout_ref.rollout.enforce_eager=${ROLLOUT_ENFORCE_EAGER} + actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} + actor_rollout_ref.rollout.n=${ROLLOUT_N} + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.n=${ROLLOUT_N_VAL} + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.enable_rollout_routing_replay=True + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} + actor_rollout_ref.rollout.max_model_len=${ROLLOUT_MAX_MODEL_LEN} + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB} + +actor_rollout_ref.rollout.engine_kwargs.vllm.kv_cache_dtype=${ROLLOUT_KV_CACHE_DTYPE} + +actor_rollout_ref.rollout.quantization=fp8 +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=${OVERLONG_BUFFER_ENABLE} + +reward.reward_kwargs.overlong_buffer_cfg.len=${OVERLONG_BUFFER_LEN} + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${OVERLONG_PENALTY_FACTOR} + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${MAX_RESPONSE_LENGTH} +) + +TRAINER=( + trainer.logger='["console","wandb"]' + trainer.project_name=${PROJECT_NAME} + trainer.experiment_name=${EXPERIMENT_NAME} + trainer.n_gpus_per_node=${NGPUS_PER_NODE} + trainer.nnodes=${NNODES} + trainer.save_freq=${SAVE_FREQ} + trainer.test_freq=${TEST_FREQ} + trainer.max_actor_ckpt_to_keep=2 + trainer.total_epochs=${TOTAL_EPOCHS} + trainer.total_training_steps=${TOTAL_TRAINING_STEPS} + trainer.resume_mode=auto + trainer.val_before_train=False + trainer.log_val_generations=100 + trainer.rollout_data_dir="${ROLLOUT_DATA_DIR}" + trainer.validation_data_dir="${VALIDATION_DATA_DIR}" + trainer.default_local_dir="${CKPTS_DIR}" +) + +EXTRA=( + actor_rollout_ref.nccl_timeout=3600 + model_engine=megatron +) + +# Simulate MXFP4 routed-expert weights while retaining the TE FP8 GEMM path. +# vLLM receives Bridge-exported weights and must not construct QAT quantizers. +QAT_ARGS=( + ++actor_rollout_ref.actor.megatron.qat.enable=True + ++actor_rollout_ref.actor.megatron.qat.mode=mxfp4_experts + ++actor_rollout_ref.actor.megatron.qat.bypass_te_fp8_assert=True + ++actor_rollout_ref.rollout.qat.enable=False +) + +########################### launch ########################### + +python3 -m verl.trainer.main_ppo \ + "${ALGORITHM[@]}" \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${ACTOR[@]}" \ + "${RECOMPUTE_ARGS[@]}" \ + "${CP_ARGS[@]}" \ + "${TIS_ARGS[@]}" \ + "${ROLLOUT[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${EXTRA[@]}" \ + "${QAT_ARGS[@]}" \ + "$@" diff --git a/tests/utils/megatron/test_hdo_checkpoint_resume.py b/tests/utils/megatron/test_hdo_checkpoint_resume.py new file mode 100644 index 00000000000..906d59e2865 --- /dev/null +++ b/tests/utils/megatron/test_hdo_checkpoint_resume.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU regression for loading HDO after verl offloads optimizer main shards. + +Use the real HDO, verl offload helpers, and engine checkpoint entry point. +The small checkpoint manager substitutes local tensor IO for distributed shards. +""" + +# ruff: noqa: E402 + +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +pytest.importorskip("megatron.core") +pytest.importorskip("transformer_engine.pytorch") + +from megatron.core import parallel_state +from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer + +from verl.utils.megatron_utils import load_megatron_optimizer, offload_megatron_optimizer +from verl.workers.engine.megatron.transformer_impl import MegatronEngine + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="HDO resume requires a GPU") + + +class CPUAdamW(torch.optim.AdamW): + """Distinct CPU type: HDO dispatches state placement via isinstance.""" + + +@pytest.fixture(scope="module", autouse=True) +def distributed(tmp_path_factory): + rendezvous = tmp_path_factory.mktemp("hdo_pg") / "init" + dist.init_process_group("nccl", init_method=rendezvous.as_uri(), rank=0, world_size=1) + parallel_state.initialize_model_parallel() + yield + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def make_optimizer(offload_fraction, overlap): + # DistributedOptimizer supplies FP32 main shards for BF16 model weights, + # plus native FP32 parameters such as DSv4's routers. + params = [torch.linspace(-0.5, 0.5, 64, device="cuda") + i for i in range(2)] + native_fp32 = torch.linspace(0.1, 0.9, 64, device="cuda") + hdo = HybridDeviceOptimizer( + params + [native_fp32], + offload_fraction=offload_fraction, + cpu_optimizer_cls=CPUAdamW, + # Use per-parameter AdamW state on both devices. TE FusedAdam's group + # step counters need additional DistributedOptimizer checkpoint IO, + # which this focused HDO/engine regression does not instantiate. + gpu_optimizer_cls=torch.optim.AdamW, + param_update_in_fp32=True, + overlap_cpu_optimizer_d2h_h2d=overlap, + lr=0.001, + weight_decay=0.01, + ) + outer = SimpleNamespace(optimizer=hdo, shard_fp32_from_float16_groups=[params]) + return outer, params + [native_fp32] + + +def step(outer, params, index): + for i, param in enumerate(params): + param.grad = torch.linspace(-0.2, 0.3, param.numel(), device="cuda") + 0.01 * (i + index) + outer.optimizer.step() + outer.optimizer.zero_grad() + torch.cuda.synchronize() + + +def save(path, outer, params): + torch.save({"params": [p.detach().cpu() for p in params], "optimizer": outer.optimizer.state_dict()}, path) + + +class TensorCheckpointManager: + def __init__(self, outer, params): + self.outer = outer + self.params = params + + def load_checkpoint(self, local_path, **kwargs): + checkpoint = torch.load(local_path, map_location="cpu", weights_only=False) + for param, saved in zip(self.params, checkpoint["params"], strict=True): + param.copy_(saved) + self.outer.optimizer.load_state_dict(checkpoint["optimizer"]) + + +def engine(outer, params, optimizer_offload=True): + # No model buffers are needed to exercise the independent optimizer shards. + return SimpleNamespace( + _is_offload_param=False, + _is_offload_optimizer=optimizer_offload, + optimizer=outer, + checkpoint_mananager=TensorCheckpointManager(outer, params), + ) + + +def assert_matches(outer, params, expected): + expected_params, expected_state = expected + for param, value in zip(params, expected_params, strict=True): + torch.testing.assert_close(param.cpu(), value, rtol=0, atol=0) + for actual, saved in zip(outer.optimizer.state_dict()["state"].values(), expected_state.values(), strict=True): + for key in ("exp_avg", "exp_avg_sq", "master_param"): + torch.testing.assert_close(actual[key].cpu(), saved[key].cpu(), rtol=0, atol=0) + assert actual["step"] == saved["step"] + + +def test_old_load_order_reproduces_missing_cpu_copy(tmp_path): + outer, params = make_optimizer(1.0, True) + step(outer, params, 1) + save(tmp_path / "checkpoint.pt", outer, params) + offload_megatron_optimizer(outer) + # Previous engine behavior: load checkpoint before restoring main shards. + TensorCheckpointManager(outer, params).load_checkpoint(tmp_path / "checkpoint.pt") + load_megatron_optimizer(outer) + with pytest.raises(KeyError): + step(outer, params, 2) + + +@pytest.mark.parametrize("offload_fraction,overlap", [(1.0, True), (0.5, True), (1.0, False)]) +@pytest.mark.parametrize("optimizer_offload", [True, False]) +def test_engine_resume_matches_uninterrupted_training(tmp_path, offload_fraction, overlap, optimizer_offload): + reference, ref_params = make_optimizer(offload_fraction, overlap) + step(reference, ref_params, 1) + step(reference, ref_params, 2) + checkpoint_path = tmp_path / "step2.pt" + save(checkpoint_path, reference, ref_params) + expected_steps = [] + for index in (3, 4): + step(reference, ref_params, index) + expected_steps.append( + ([p.detach().cpu().clone() for p in ref_params], deepcopy(reference.optimizer.state_dict()["state"])) + ) + + # Resume twice, including a save after the first restored step. + for index, expected in zip((3, 4), expected_steps, strict=True): + restored, params = make_optimizer(offload_fraction, overlap) + if optimizer_offload: + offload_megatron_optimizer(restored) + MegatronEngine.load_checkpoint(engine(restored, params, optimizer_offload), str(checkpoint_path)) + load_megatron_optimizer(restored) + step(restored, params, index) + assert_matches(restored, params, expected) + checkpoint_path = tmp_path / f"step{index}.pt" + save(checkpoint_path, restored, params) diff --git a/tests/utils/modelopt/test_mxfp4_checkpoint.py b/tests/utils/modelopt/test_mxfp4_checkpoint.py new file mode 100644 index 00000000000..537d2633e1c --- /dev/null +++ b/tests/utils/modelopt/test_mxfp4_checkpoint.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU regressions for MXFP4 QAT with native Megatron/TE checkpoint IO. + +Run with the training stack (ModelOpt 0.44, TE >= 2, Megatron-Core) on one GPU: + pytest -q tests/utils/modelopt/test_mxfp4_checkpoint.py +""" + +# ruff: noqa: E402 + +import pickle +from contextlib import nullcontext +from importlib.metadata import version + +import pytest +import torch +import torch.distributed as dist + +pytest.importorskip("megatron.core") +mtq = pytest.importorskip("modelopt.torch.quantization") +te = pytest.importorskip("transformer_engine.pytorch") + +from megatron.core import dist_checkpointing, parallel_state +from megatron.core.extensions.transformer_engine import ( + TEColumnParallelGroupedLinear, + TEColumnParallelLinear, + TERowParallelGroupedLinear, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.experts import GroupedMLPSubmodules, TEGroupedMLP +from megatron.core.transformer.transformer_config import TransformerConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from transformer_engine.common.recipe import DelayedScaling, Float8BlockScaling, Format + +from verl.utils.modelopt.checkpoint import preserve_mxfp4_checkpoint_methods +from verl.utils.modelopt.quantize import apply_qat, build_quantize_config + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron TE checkpoint regression needs a GPU") + + +@pytest.fixture(scope="module", autouse=True) +def distributed(tmp_path_factory): + rendezvous = tmp_path_factory.mktemp("mxfp4_pg") / "init" + dist.init_process_group("nccl", init_method=rendezvous.as_uri(), rank=0, world_size=1) + parallel_state.initialize_model_parallel() + model_parallel_cuda_manual_seed(1234) + yield + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +@pytest.fixture(autouse=True) +def enable_te_fp8(monkeypatch): + import modelopt.torch.quantization.plugins.transformer_engine as modelopt_te + + monkeypatch.setattr(modelopt_te, "_assert_te_fp8_enabled", lambda: None) + + +def make_model(recipe="blockwise"): + config = TransformerConfig( + num_layers=1, + hidden_size=256, + num_attention_heads=4, + ffn_hidden_size=256, + moe_ffn_hidden_size=256, + num_moe_experts=2, + moe_grouped_gemm=True, + gated_linear_unit=True, + activation_func=torch.nn.functional.silu, + add_bias_linear=False, + gradient_accumulation_fusion=False, + params_dtype=torch.bfloat16, + bf16=True, + fp8="hybrid" if recipe else None, + fp8_recipe=recipe or "delayed", + ) + model = MegatronModule(config) + model.layer = MegatronModule(config) + model.layer.mlp = MegatronModule(config) + model.layer.mlp.experts = TEGroupedMLP( + num_local_experts=2, + config=config, + submodules=GroupedMLPSubmodules( + linear_fc1=TEColumnParallelGroupedLinear, + linear_fc2=TERowParallelGroupedLinear, + ), + pg_collection=ProcessGroupCollection.use_mpu_process_groups(), + ) + # ModelOpt also converts non-expert linears, even though their quantizers + # are disabled. Their TE state and ordinary weights must survive too. + model.layer.dense = TEColumnParallelLinear( + 256, + 256, + config=config, + init_method=config.init_method, + gather_output=False, + bias=False, + skip_bias_add=True, + is_expert=False, + ) + model.layer.plain = torch.nn.Linear(256, 256, bias=False, device="cuda", dtype=torch.bfloat16) + return model + + +def train_step(model, optimizer, recipe): + experts = model.layer.mlp.experts + for linear in (experts.linear_fc1, experts.linear_fc2, model.layer.dense): + linear.is_first_microbatch = True + fp8_recipe = ( + Float8BlockScaling(fp8_format=Format.HYBRID) + if recipe == "blockwise" + else DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=4) + ) + context = te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe) if recipe else nullcontext() + optimizer.zero_grad(set_to_none=True) + inputs = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with context: + output, _ = experts(inputs, torch.tensor([128, 128]), torch.ones(256, device="cuda")) + output, _ = model.layer.dense(output) + output = model.layer.plain(output) + # Delayed scaling starts with unit scales. A mean over 65,536 values + # underflows its first FP8 backward; a sum keeps this IO test in range. + loss = output.float().square().sum() + loss.backward() + assert torch.isfinite(loss) + assert any(p.grad is not None and p.grad.abs().max() > 0 for p in model.parameters()) + optimizer.step() + + +def sharded(model): + return {"model": model.sharded_state_dict(metadata={"dp_cp_group": parallel_state.get_data_parallel_group()})} + + +def test_modelopt_044_reproduces_fp8_save_error(): + if not version("nvidia-modelopt").startswith("0.44."): + pytest.skip("Upstream failure is specific to ModelOpt 0.44") + model = make_model() + mtq.quantize(model, build_quantize_config("mxfp4_experts")) + train_step(model, torch.optim.Adam(model.parameters(), lr=1e-3), "blockwise") + with pytest.raises(pickle.UnpicklingError, match="persistent"): + sharded(model) + + +@pytest.mark.parametrize("recipe", ["blockwise", "delayed", None]) +def test_mxfp4_save_load_and_continue(tmp_path, recipe): + model = make_model(recipe) + original_parameter_names = list(dict(model.named_parameters())) + original_shard_keys = set(sharded(model)["model"]) + original_heterogeneous = model.config.hetereogenous_dist_checkpoint + apply_qat(model, "mxfp4_experts") + assert list(dict(model.named_parameters())) == original_parameter_names + assert model.config.hetereogenous_dist_checkpoint == original_heterogeneous + assert set(sharded(model)["model"]) == original_shard_keys + assert not model.layer.dense.weight_quantizer.is_enabled + assert not model.layer.plain.weight_quantizer.is_enabled + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + for step in range(2): + train_step(model, optimizer, recipe) + save_dir = tmp_path / f"step_{step}" + save_dir.mkdir() + dist_checkpointing.save(sharded(model), str(save_dir)) + expected_weights = {name: p.detach().clone() for name, p in model.named_parameters()} + optimizer_path = tmp_path / "optimizer.pt" + torch.save(optimizer.state_dict(), optimizer_path) + + restored = make_model(recipe) + apply_qat(restored, "mxfp4_experts") + loaded = dist_checkpointing.load(sharded(restored), str(save_dir)) + restored.load_state_dict(loaded["model"], strict=True) + for name, p in restored.named_parameters(): + torch.testing.assert_close(p, expected_weights[name], rtol=0, atol=0) + for source, target in zip( + (model.layer.mlp.experts.linear_fc1, model.layer.mlp.experts.linear_fc2, model.layer.dense), + (restored.layer.mlp.experts.linear_fc1, restored.layer.mlp.experts.linear_fc2, restored.layer.dense), + strict=True, + ): + assert target.weight_quantizer.is_enabled == source.weight_quantizer.is_enabled + assert not target.input_quantizer.is_enabled + source_extra = source.get_extra_state() + target_extra = target.get_extra_state() + source_state = pickle.loads(source_extra.cpu().numpy().tobytes()) if source_extra.numel() else None + target_state = pickle.loads(target_extra.cpu().numpy().tobytes()) if target_extra.numel() else None + if recipe: + if hasattr(target, "num_gemms"): + assert target_state["extra_fp8_variables"]["num_gemms"] == 2 + assert type(target_state["recipe"]) is type(source_state["recipe"]) + if recipe == "delayed": + for name in ("scale_fwd", "scale_bwd", "amax_history_fwd", "amax_history_bwd"): + torch.testing.assert_close(target_state[name], source_state[name], rtol=0, atol=0) + else: + assert target_state is None + restored_optimizer = torch.optim.Adam(restored.parameters(), lr=1e-3) + restored_optimizer.load_state_dict(torch.load(optimizer_path, weights_only=True)) + train_step(restored, restored_optimizer, recipe) + assert any(not torch.equal(p, expected_weights[name]) for name, p in restored.named_parameters()) + + +def test_persistent_quantizer_state_is_rejected(): + model = torch.nn.Module() + with pytest.raises(RuntimeError, match="requires stateless quantizers"): + with preserve_mxfp4_checkpoint_methods(model): + model.weight_quantizer = TensorQuantizer() + model.weight_quantizer.register_buffer("_amax", torch.tensor(1.0)) + + +def test_other_qat_modes_keep_modelopt_checkpoint_methods(monkeypatch): + model = torch.nn.Linear(32, 32) + modelopt_method = lambda: "modelopt state" + + def quantize(module, config): + module.get_extra_state = modelopt_method + + monkeypatch.setattr(mtq, "quantize", quantize) + apply_qat(model, "w4a16") + assert model.get_extra_state is modelopt_method diff --git a/verl/utils/modelopt/checkpoint.py b/verl/utils/modelopt/checkpoint.py new file mode 100644 index 00000000000..1319705bdcf --- /dev/null +++ b/verl/utils/modelopt/checkpoint.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Keep native Megatron checkpoints for stateless MXFP4 expert QAT.""" + +from contextlib import contextmanager + + +@contextmanager +def preserve_mxfp4_checkpoint_methods(model): + """Preserve checkpoint methods across ModelOpt's in-place module conversion. + + ModelOpt 0.44 wraps TE extra state in a torch.save byte tensor, whereas + Megatron's grouped-linear split/merge expects TE's native pickle payload. + Its grouped-linear loader also discards the per-GEMM TE extra states. + Dynamic MXFP4 expert QAT has no persistent quantizer tensors: its config + is reapplied before checkpoint loading, so native checkpointing suffices. + Keep the QAT forward classes and quantizers; restore only checkpoint IO. + """ + method_names = ("get_extra_state", "set_extra_state", "sharded_state_dict", "_load_from_state_dict") + original_methods = [ + (module, {name: getattr(module, name) for name in method_names if callable(getattr(module, name, None))}) + for module in model.modules() + ] + # ModelOpt enables per-layer checkpoint keys for quantizer metadata. This + # mode does not need them; preserve compatibility with non-QAT checkpoints. + original_configs = {} + for module, _ in original_methods: + config = getattr(module, "config", None) + if hasattr(config, "hetereogenous_dist_checkpoint"): + original_configs[id(config)] = (config, config.hetereogenous_dist_checkpoint) + + yield + + from modelopt.torch.quantization.nn import TensorQuantizer + + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and module.state_dict(): + raise RuntimeError(f"MXFP4 native checkpointing requires stateless quantizers; {name} has persistent state") + for module, methods in original_methods: + for name, method in methods.items(): + if getattr(module, name) != method: + setattr(module, name, method) + for config, heterogeneous in original_configs.values(): + config.hetereogenous_dist_checkpoint = heterogeneous diff --git a/verl/utils/modelopt/quantize.py b/verl/utils/modelopt/quantize.py index b16836f8209..c2313b30851 100644 --- a/verl/utils/modelopt/quantize.py +++ b/verl/utils/modelopt/quantize.py @@ -100,5 +100,11 @@ def apply_qat( ) -> nn.Module: """Apply Quantization-Aware Training to a Megatron model.""" config = build_quantize_config(qat_mode, ignore_patterns) - mtq.quantize(model, config) + if qat_mode == "mxfp4_experts": + from verl.utils.modelopt.checkpoint import preserve_mxfp4_checkpoint_methods + + with preserve_mxfp4_checkpoint_methods(model): + mtq.quantize(model, config) + else: + mtq.quantize(model, config) return model diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index b2f287cb3fa..a5fc9198657 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -595,8 +595,7 @@ def initialize(self): if hasattr(_mo_te, "_assert_te_fp8_enabled"): _mo_te._assert_te_fp8_enabled = lambda: None logger.warning( - "QAT: bypassing modelopt _assert_te_fp8_enabled " - "(weight-only fake-quant + TE FP8 GEMM)" + "QAT: bypassing modelopt _assert_te_fp8_enabled (weight-only fake-quant + TE FP8 GEMM)" ) from verl.utils.modelopt import apply_qat_to_modules @@ -828,6 +827,12 @@ def load_checkpoint( """ if self._is_offload_param: load_megatron_model_to_gpu(self.module) + if self._is_offload_optimizer: + # HDO rebuilds its CPU/GPU parameter mappings in load_state_dict. + # Restore the optimizer's main parameter shards first: loading only + # the model leaves the offloaded copies on CPU, where HDO treats + # them as native CPU parameters and loses the GPU copy mappings. + load_megatron_optimizer(self.optimizer) self.checkpoint_mananager.load_checkpoint( local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load ) From a8d21f9090e01412c8334a5ee4d7fc2f2face81f Mon Sep 17 00:00:00 2001 From: Sophia Yang Date: Mon, 14 Sep 2026 09:25:01 -0700 Subject: [PATCH 4/6] [test] remove MXFP4 checkpoint and HDO resume regression files Remove the two checkpoint regression test files from this branch as requested. Co-authored-by: OpenAI Codex --- .../megatron/test_hdo_checkpoint_resume.py | 159 ------------ tests/utils/modelopt/test_mxfp4_checkpoint.py | 229 ------------------ 2 files changed, 388 deletions(-) delete mode 100644 tests/utils/megatron/test_hdo_checkpoint_resume.py delete mode 100644 tests/utils/modelopt/test_mxfp4_checkpoint.py diff --git a/tests/utils/megatron/test_hdo_checkpoint_resume.py b/tests/utils/megatron/test_hdo_checkpoint_resume.py deleted file mode 100644 index 906d59e2865..00000000000 --- a/tests/utils/megatron/test_hdo_checkpoint_resume.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""GPU regression for loading HDO after verl offloads optimizer main shards. - -Use the real HDO, verl offload helpers, and engine checkpoint entry point. -The small checkpoint manager substitutes local tensor IO for distributed shards. -""" - -# ruff: noqa: E402 - -from copy import deepcopy -from types import SimpleNamespace - -import pytest -import torch -import torch.distributed as dist - -pytest.importorskip("megatron.core") -pytest.importorskip("transformer_engine.pytorch") - -from megatron.core import parallel_state -from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer - -from verl.utils.megatron_utils import load_megatron_optimizer, offload_megatron_optimizer -from verl.workers.engine.megatron.transformer_impl import MegatronEngine - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="HDO resume requires a GPU") - - -class CPUAdamW(torch.optim.AdamW): - """Distinct CPU type: HDO dispatches state placement via isinstance.""" - - -@pytest.fixture(scope="module", autouse=True) -def distributed(tmp_path_factory): - rendezvous = tmp_path_factory.mktemp("hdo_pg") / "init" - dist.init_process_group("nccl", init_method=rendezvous.as_uri(), rank=0, world_size=1) - parallel_state.initialize_model_parallel() - yield - parallel_state.destroy_model_parallel() - dist.destroy_process_group() - - -def make_optimizer(offload_fraction, overlap): - # DistributedOptimizer supplies FP32 main shards for BF16 model weights, - # plus native FP32 parameters such as DSv4's routers. - params = [torch.linspace(-0.5, 0.5, 64, device="cuda") + i for i in range(2)] - native_fp32 = torch.linspace(0.1, 0.9, 64, device="cuda") - hdo = HybridDeviceOptimizer( - params + [native_fp32], - offload_fraction=offload_fraction, - cpu_optimizer_cls=CPUAdamW, - # Use per-parameter AdamW state on both devices. TE FusedAdam's group - # step counters need additional DistributedOptimizer checkpoint IO, - # which this focused HDO/engine regression does not instantiate. - gpu_optimizer_cls=torch.optim.AdamW, - param_update_in_fp32=True, - overlap_cpu_optimizer_d2h_h2d=overlap, - lr=0.001, - weight_decay=0.01, - ) - outer = SimpleNamespace(optimizer=hdo, shard_fp32_from_float16_groups=[params]) - return outer, params + [native_fp32] - - -def step(outer, params, index): - for i, param in enumerate(params): - param.grad = torch.linspace(-0.2, 0.3, param.numel(), device="cuda") + 0.01 * (i + index) - outer.optimizer.step() - outer.optimizer.zero_grad() - torch.cuda.synchronize() - - -def save(path, outer, params): - torch.save({"params": [p.detach().cpu() for p in params], "optimizer": outer.optimizer.state_dict()}, path) - - -class TensorCheckpointManager: - def __init__(self, outer, params): - self.outer = outer - self.params = params - - def load_checkpoint(self, local_path, **kwargs): - checkpoint = torch.load(local_path, map_location="cpu", weights_only=False) - for param, saved in zip(self.params, checkpoint["params"], strict=True): - param.copy_(saved) - self.outer.optimizer.load_state_dict(checkpoint["optimizer"]) - - -def engine(outer, params, optimizer_offload=True): - # No model buffers are needed to exercise the independent optimizer shards. - return SimpleNamespace( - _is_offload_param=False, - _is_offload_optimizer=optimizer_offload, - optimizer=outer, - checkpoint_mananager=TensorCheckpointManager(outer, params), - ) - - -def assert_matches(outer, params, expected): - expected_params, expected_state = expected - for param, value in zip(params, expected_params, strict=True): - torch.testing.assert_close(param.cpu(), value, rtol=0, atol=0) - for actual, saved in zip(outer.optimizer.state_dict()["state"].values(), expected_state.values(), strict=True): - for key in ("exp_avg", "exp_avg_sq", "master_param"): - torch.testing.assert_close(actual[key].cpu(), saved[key].cpu(), rtol=0, atol=0) - assert actual["step"] == saved["step"] - - -def test_old_load_order_reproduces_missing_cpu_copy(tmp_path): - outer, params = make_optimizer(1.0, True) - step(outer, params, 1) - save(tmp_path / "checkpoint.pt", outer, params) - offload_megatron_optimizer(outer) - # Previous engine behavior: load checkpoint before restoring main shards. - TensorCheckpointManager(outer, params).load_checkpoint(tmp_path / "checkpoint.pt") - load_megatron_optimizer(outer) - with pytest.raises(KeyError): - step(outer, params, 2) - - -@pytest.mark.parametrize("offload_fraction,overlap", [(1.0, True), (0.5, True), (1.0, False)]) -@pytest.mark.parametrize("optimizer_offload", [True, False]) -def test_engine_resume_matches_uninterrupted_training(tmp_path, offload_fraction, overlap, optimizer_offload): - reference, ref_params = make_optimizer(offload_fraction, overlap) - step(reference, ref_params, 1) - step(reference, ref_params, 2) - checkpoint_path = tmp_path / "step2.pt" - save(checkpoint_path, reference, ref_params) - expected_steps = [] - for index in (3, 4): - step(reference, ref_params, index) - expected_steps.append( - ([p.detach().cpu().clone() for p in ref_params], deepcopy(reference.optimizer.state_dict()["state"])) - ) - - # Resume twice, including a save after the first restored step. - for index, expected in zip((3, 4), expected_steps, strict=True): - restored, params = make_optimizer(offload_fraction, overlap) - if optimizer_offload: - offload_megatron_optimizer(restored) - MegatronEngine.load_checkpoint(engine(restored, params, optimizer_offload), str(checkpoint_path)) - load_megatron_optimizer(restored) - step(restored, params, index) - assert_matches(restored, params, expected) - checkpoint_path = tmp_path / f"step{index}.pt" - save(checkpoint_path, restored, params) diff --git a/tests/utils/modelopt/test_mxfp4_checkpoint.py b/tests/utils/modelopt/test_mxfp4_checkpoint.py deleted file mode 100644 index 537d2633e1c..00000000000 --- a/tests/utils/modelopt/test_mxfp4_checkpoint.py +++ /dev/null @@ -1,229 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""GPU regressions for MXFP4 QAT with native Megatron/TE checkpoint IO. - -Run with the training stack (ModelOpt 0.44, TE >= 2, Megatron-Core) on one GPU: - pytest -q tests/utils/modelopt/test_mxfp4_checkpoint.py -""" - -# ruff: noqa: E402 - -import pickle -from contextlib import nullcontext -from importlib.metadata import version - -import pytest -import torch -import torch.distributed as dist - -pytest.importorskip("megatron.core") -mtq = pytest.importorskip("modelopt.torch.quantization") -te = pytest.importorskip("transformer_engine.pytorch") - -from megatron.core import dist_checkpointing, parallel_state -from megatron.core.extensions.transformer_engine import ( - TEColumnParallelGroupedLinear, - TEColumnParallelLinear, - TERowParallelGroupedLinear, -) -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.module import MegatronModule -from megatron.core.transformer.moe.experts import GroupedMLPSubmodules, TEGroupedMLP -from megatron.core.transformer.transformer_config import TransformerConfig -from modelopt.torch.quantization.nn import TensorQuantizer -from transformer_engine.common.recipe import DelayedScaling, Float8BlockScaling, Format - -from verl.utils.modelopt.checkpoint import preserve_mxfp4_checkpoint_methods -from verl.utils.modelopt.quantize import apply_qat, build_quantize_config - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron TE checkpoint regression needs a GPU") - - -@pytest.fixture(scope="module", autouse=True) -def distributed(tmp_path_factory): - rendezvous = tmp_path_factory.mktemp("mxfp4_pg") / "init" - dist.init_process_group("nccl", init_method=rendezvous.as_uri(), rank=0, world_size=1) - parallel_state.initialize_model_parallel() - model_parallel_cuda_manual_seed(1234) - yield - parallel_state.destroy_model_parallel() - dist.destroy_process_group() - - -@pytest.fixture(autouse=True) -def enable_te_fp8(monkeypatch): - import modelopt.torch.quantization.plugins.transformer_engine as modelopt_te - - monkeypatch.setattr(modelopt_te, "_assert_te_fp8_enabled", lambda: None) - - -def make_model(recipe="blockwise"): - config = TransformerConfig( - num_layers=1, - hidden_size=256, - num_attention_heads=4, - ffn_hidden_size=256, - moe_ffn_hidden_size=256, - num_moe_experts=2, - moe_grouped_gemm=True, - gated_linear_unit=True, - activation_func=torch.nn.functional.silu, - add_bias_linear=False, - gradient_accumulation_fusion=False, - params_dtype=torch.bfloat16, - bf16=True, - fp8="hybrid" if recipe else None, - fp8_recipe=recipe or "delayed", - ) - model = MegatronModule(config) - model.layer = MegatronModule(config) - model.layer.mlp = MegatronModule(config) - model.layer.mlp.experts = TEGroupedMLP( - num_local_experts=2, - config=config, - submodules=GroupedMLPSubmodules( - linear_fc1=TEColumnParallelGroupedLinear, - linear_fc2=TERowParallelGroupedLinear, - ), - pg_collection=ProcessGroupCollection.use_mpu_process_groups(), - ) - # ModelOpt also converts non-expert linears, even though their quantizers - # are disabled. Their TE state and ordinary weights must survive too. - model.layer.dense = TEColumnParallelLinear( - 256, - 256, - config=config, - init_method=config.init_method, - gather_output=False, - bias=False, - skip_bias_add=True, - is_expert=False, - ) - model.layer.plain = torch.nn.Linear(256, 256, bias=False, device="cuda", dtype=torch.bfloat16) - return model - - -def train_step(model, optimizer, recipe): - experts = model.layer.mlp.experts - for linear in (experts.linear_fc1, experts.linear_fc2, model.layer.dense): - linear.is_first_microbatch = True - fp8_recipe = ( - Float8BlockScaling(fp8_format=Format.HYBRID) - if recipe == "blockwise" - else DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=4) - ) - context = te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe) if recipe else nullcontext() - optimizer.zero_grad(set_to_none=True) - inputs = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) - with context: - output, _ = experts(inputs, torch.tensor([128, 128]), torch.ones(256, device="cuda")) - output, _ = model.layer.dense(output) - output = model.layer.plain(output) - # Delayed scaling starts with unit scales. A mean over 65,536 values - # underflows its first FP8 backward; a sum keeps this IO test in range. - loss = output.float().square().sum() - loss.backward() - assert torch.isfinite(loss) - assert any(p.grad is not None and p.grad.abs().max() > 0 for p in model.parameters()) - optimizer.step() - - -def sharded(model): - return {"model": model.sharded_state_dict(metadata={"dp_cp_group": parallel_state.get_data_parallel_group()})} - - -def test_modelopt_044_reproduces_fp8_save_error(): - if not version("nvidia-modelopt").startswith("0.44."): - pytest.skip("Upstream failure is specific to ModelOpt 0.44") - model = make_model() - mtq.quantize(model, build_quantize_config("mxfp4_experts")) - train_step(model, torch.optim.Adam(model.parameters(), lr=1e-3), "blockwise") - with pytest.raises(pickle.UnpicklingError, match="persistent"): - sharded(model) - - -@pytest.mark.parametrize("recipe", ["blockwise", "delayed", None]) -def test_mxfp4_save_load_and_continue(tmp_path, recipe): - model = make_model(recipe) - original_parameter_names = list(dict(model.named_parameters())) - original_shard_keys = set(sharded(model)["model"]) - original_heterogeneous = model.config.hetereogenous_dist_checkpoint - apply_qat(model, "mxfp4_experts") - assert list(dict(model.named_parameters())) == original_parameter_names - assert model.config.hetereogenous_dist_checkpoint == original_heterogeneous - assert set(sharded(model)["model"]) == original_shard_keys - assert not model.layer.dense.weight_quantizer.is_enabled - assert not model.layer.plain.weight_quantizer.is_enabled - optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) - for step in range(2): - train_step(model, optimizer, recipe) - save_dir = tmp_path / f"step_{step}" - save_dir.mkdir() - dist_checkpointing.save(sharded(model), str(save_dir)) - expected_weights = {name: p.detach().clone() for name, p in model.named_parameters()} - optimizer_path = tmp_path / "optimizer.pt" - torch.save(optimizer.state_dict(), optimizer_path) - - restored = make_model(recipe) - apply_qat(restored, "mxfp4_experts") - loaded = dist_checkpointing.load(sharded(restored), str(save_dir)) - restored.load_state_dict(loaded["model"], strict=True) - for name, p in restored.named_parameters(): - torch.testing.assert_close(p, expected_weights[name], rtol=0, atol=0) - for source, target in zip( - (model.layer.mlp.experts.linear_fc1, model.layer.mlp.experts.linear_fc2, model.layer.dense), - (restored.layer.mlp.experts.linear_fc1, restored.layer.mlp.experts.linear_fc2, restored.layer.dense), - strict=True, - ): - assert target.weight_quantizer.is_enabled == source.weight_quantizer.is_enabled - assert not target.input_quantizer.is_enabled - source_extra = source.get_extra_state() - target_extra = target.get_extra_state() - source_state = pickle.loads(source_extra.cpu().numpy().tobytes()) if source_extra.numel() else None - target_state = pickle.loads(target_extra.cpu().numpy().tobytes()) if target_extra.numel() else None - if recipe: - if hasattr(target, "num_gemms"): - assert target_state["extra_fp8_variables"]["num_gemms"] == 2 - assert type(target_state["recipe"]) is type(source_state["recipe"]) - if recipe == "delayed": - for name in ("scale_fwd", "scale_bwd", "amax_history_fwd", "amax_history_bwd"): - torch.testing.assert_close(target_state[name], source_state[name], rtol=0, atol=0) - else: - assert target_state is None - restored_optimizer = torch.optim.Adam(restored.parameters(), lr=1e-3) - restored_optimizer.load_state_dict(torch.load(optimizer_path, weights_only=True)) - train_step(restored, restored_optimizer, recipe) - assert any(not torch.equal(p, expected_weights[name]) for name, p in restored.named_parameters()) - - -def test_persistent_quantizer_state_is_rejected(): - model = torch.nn.Module() - with pytest.raises(RuntimeError, match="requires stateless quantizers"): - with preserve_mxfp4_checkpoint_methods(model): - model.weight_quantizer = TensorQuantizer() - model.weight_quantizer.register_buffer("_amax", torch.tensor(1.0)) - - -def test_other_qat_modes_keep_modelopt_checkpoint_methods(monkeypatch): - model = torch.nn.Linear(32, 32) - modelopt_method = lambda: "modelopt state" - - def quantize(module, config): - module.get_extra_state = modelopt_method - - monkeypatch.setattr(mtq, "quantize", quantize) - apply_qat(model, "w4a16") - assert model.get_extra_state is modelopt_method From d816d5cb465c05da6801607eefe2ceb02839ea9b Mon Sep 17 00:00:00 2001 From: Sophia Yang Date: Mon, 14 Sep 2026 09:30:05 -0700 Subject: [PATCH 5/6] [example] document Megatron prerequisite in DSv4 QAT recipe Record the HDO FP32-shard detach prerequisite, upstream fix and validation provenance in the standalone MXFP4 QAT script. Remove the length-only recipe and its stale reference from the QAT header. Validation: bash -n and git diff --check pass; the executable QAT script body is byte-for-byte unchanged. Co-authored-by: OpenAI Codex --- ..._deepseek_v4_flash_length_4096_megatron.sh | 287 ------------------ ...pseek_v4_flash_length_4096_qat_megatron.sh | 37 ++- 2 files changed, 36 insertions(+), 288 deletions(-) delete mode 100755 examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh diff --git a/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh deleted file mode 100755 index 902a42b4316..00000000000 --- a/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_megatron.sh +++ /dev/null @@ -1,287 +0,0 @@ -#!/usr/bin/env bash -# DeepSeek-V4-Flash-0731: response-length experiment, 40 steps, validation every 10. -# Standalone length-only baseline; the MXFP4 QAT arm is a separate script. -# Plan A: prompt/response 1024/4096, overlong buffer 512 with penalty 1.0, -# clip high 0.28 / dual clip 10.0, no thinking, no QAT; retain TE FP8 training. -# Set MODEL_PATH, TRAIN_FILE and TEST_FILE in the launch environment. -# Optional: set OUTPUT_DIR or override each checkpoint/rollout/validation directory. -# For two segments, use TOTAL_TRAINING_STEPS=20, then 40 with the same checkpoint directory. -# Requires DeepSeek-V4 support in Megatron-Bridge, Megatron-Core and vLLM. - -set -euo pipefail - -export CUDA_DEVICE_MAX_CONNECTIONS=1 -export NCCL_NVLS_ENABLE=0 -export TORCH_NCCL_AVOID_RECORD_STREAMS=1 -export VLLM_USE_V1=1 -export PYTHONUNBUFFERED=1 -export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} -export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 # Use FP32 scales for TE block-wise FP8 -export VLLM_USE_DEEP_GEMM_E8M0=0 # Use FP32 scales for vLLM with DeepGEMM - -############################### configs ################################ - -: "${MODEL_PATH:?Set MODEL_PATH to the model checkpoint directory}" -: "${TRAIN_FILE:?Set TRAIN_FILE to the training parquet file}" -: "${TEST_FILE:?Set TEST_FILE to the validation parquet file}" -OUTPUT_DIR=${OUTPUT_DIR:-./outputs} -NNODES=${NNODES:-16} -NGPUS_PER_NODE=${NGPUS_PER_NODE:-8} -TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} -PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-32} -ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU:-1} -DATALOADER_NUM_WORKERS=${DATALOADER_NUM_WORKERS:-8} -MAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-1024} -MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-4096} -PPO_MAX_TOKEN_LEN_PER_GPU=${PPO_MAX_TOKEN_LEN_PER_GPU:-5120} - -USE_KL_IN_REWARD=${USE_KL_IN_REWARD:-False} -KL_COEF=${KL_COEF:-0.001} -CLIP_RATIO_LOW=${CLIP_RATIO_LOW:-0.2} -CLIP_RATIO_HIGH=${CLIP_RATIO_HIGH:-0.28} -CLIP_RATIO_C=${CLIP_RATIO_C:-10.0} -ACTOR_LR=${ACTOR_LR:-1e-6} -OPTIMIZER_OFFLOAD_FRACTION=${OPTIMIZER_OFFLOAD_FRACTION:-1.0} - -ACTOR_TP=${ACTOR_TP:-1} -ACTOR_PP=${ACTOR_PP:-8} -ACTOR_VPP=${ACTOR_VPP:-null} -ACTOR_EP=${ACTOR_EP:-16} -ACTOR_ETP=${ACTOR_ETP:-1} -ACTOR_CP=${ACTOR_CP:-1} -PIPELINE_MODEL_PARALLEL_LAYOUT=${PIPELINE_MODEL_PARALLEL_LAYOUT:-"Et*6|t*6|t*6|t*5|t*5|t*5|t*5|t*5L"} -RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} -ALL_OFFLOAD=${ALL_OFFLOAD:-True} - -ROLLOUT_EP=${ROLLOUT_EP:-8} -ROLLOUT_N=${ROLLOUT_N:-8} -ROLLOUT_N_VAL=${ROLLOUT_N_VAL:-16} -ROLLOUT_ENFORCE_EAGER=${ROLLOUT_ENFORCE_EAGER:-False} -ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.4} -ROLLOUT_MAX_MODEL_LEN=${ROLLOUT_MAX_MODEL_LEN:-5120} -# Keep the proven prefill budget independent of the longer response/context cap. -ROLLOUT_MAX_NUM_BATCHED_TOKENS=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-3072} -ROLLOUT_KV_CACHE_DTYPE=${ROLLOUT_KV_CACHE_DTYPE:-fp8} -ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB:-512} -ROUTER_REPLAY_MODE=${ROUTER_REPLAY_MODE:-R3} -ROLLOUT_IS_MODE=${ROLLOUT_IS_MODE:-token} -ROLLOUT_IS_THRESHOLD=${ROLLOUT_IS_THRESHOLD:-2.0} - -TOTAL_EPOCHS=${TOTAL_EPOCHS:-10} -# The submitter sets this to 20 for segment 1 and 40 for the resumed segment. -TOTAL_TRAINING_STEPS=${TOTAL_TRAINING_STEPS:-40} -SAVE_FREQ=${SAVE_FREQ:-20} -TEST_FREQ=${TEST_FREQ:-10} -PROJECT_NAME=${PROJECT_NAME:-verl_deepseek_v4_flash} -EXPERIMENT_NAME=${EXPERIMENT_NAME:-dsv4-0731-noqat-len4096-40step-val10} -CKPTS_DIR=${CKPTS_DIR:-$OUTPUT_DIR/checkpoints/$PROJECT_NAME/$EXPERIMENT_NAME} -ROLLOUT_DATA_DIR=${ROLLOUT_DATA_DIR:-$OUTPUT_DIR/rollout_dumps/$EXPERIMENT_NAME} -VALIDATION_DATA_DIR=${VALIDATION_DATA_DIR:-$OUTPUT_DIR/validation_dumps/$EXPERIMENT_NAME} -OVERLONG_BUFFER_LEN=${OVERLONG_BUFFER_LEN:-512} -OVERLONG_BUFFER_ENABLE=${OVERLONG_BUFFER_ENABLE:-True} -OVERLONG_PENALTY_FACTOR=${OVERLONG_PENALTY_FACTOR:-1.0} -ENABLE_THINKING=${ENABLE_THINKING:-False} - -########################### parameter arrays ########################### - -ALGORITHM=( - algorithm.adv_estimator=grpo - algorithm.use_kl_in_reward=${USE_KL_IN_REWARD} - algorithm.kl_ctrl.kl_coef=${KL_COEF} -) - -DATA=( - data.train_files="$TRAIN_FILE" - data.val_files="$TEST_FILE" - data.train_batch_size=${TRAIN_BATCH_SIZE} - data.prompt_key=prompt - data.return_raw_chat=True - data.max_prompt_length=${MAX_PROMPT_LENGTH} - data.max_response_length=${MAX_RESPONSE_LENGTH} - data.filter_overlong_prompts=False - data.filter_overlong_prompts_workers=64 - data.truncation=error - data.dataloader_num_workers=${DATALOADER_NUM_WORKERS} - +data.apply_chat_template_kwargs.enable_thinking=${ENABLE_THINKING} -) - -MODEL=( - actor_rollout_ref.model.path="$MODEL_PATH" - actor_rollout_ref.model.trust_remote_code=True - actor_rollout_ref.model.use_fused_kernels=True - actor_rollout_ref.model.use_remove_padding=True - actor_rollout_ref.model.enable_gradient_checkpointing=True -) - -ACTOR=( - actor_rollout_ref.actor.optim.lr=${ACTOR_LR} - +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} - +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True - +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True - # actor_rollout_ref.actor.policy_loss.loss_mode=${LOSS_MODE} - # actor_rollout_ref.actor.loss_agg_mode=${LOSS_AGG_MODE} - # actor_rollout_ref.actor.use_kl_loss=${USE_KL_LOSS} - # actor_rollout_ref.actor.kl_loss_coef=${KL_LOSS_COEF} - actor_rollout_ref.actor.clip_ratio_low=${CLIP_RATIO_LOW} - actor_rollout_ref.actor.clip_ratio_high=${CLIP_RATIO_HIGH} - actor_rollout_ref.actor.clip_ratio_c=${CLIP_RATIO_C} - actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} - actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU} - actor_rollout_ref.actor.use_dynamic_bsz=True - actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} - actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} - actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} - actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} - actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} - actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} - actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} - actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD} - actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD} - # grad_offload removed upstream by verl 559c337a (#7544, BREAKING) as unused -- - # setting it now raises ConfigAttributeError: Key 'grad_offload' is not in struct. - actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTER_REPLAY_MODE} - ++actor_rollout_ref.actor.megatron.override_transformer_config.apply_dsa_kernel_fusion=True - ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_use_sparse_loss=True - ++actor_rollout_ref.actor.megatron.override_transformer_config.dsa_indexer_loss_coeff=0.0 - ++actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_mhc=False - # moe_router_fusion: TE fused-router backward conflicts with router_replay R3 (inplace-modified - # routing map in autograd) — RuntimeError at te/pytorch/router.py:72. Keep OFF with R3. - # ++actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_fusion=True - # ++actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_padding_for_quantization=True - +actor_rollout_ref.actor.megatron.override_transformer_config.fp8="hybrid" - +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe="blockwise" - +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe="blockwise" - # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_param=True - # +actor_rollout_ref.actor.megatron.override_ddp_config.fp8_param_gather=True - # +actor_rollout_ref.actor.megatron.override_ddp_config.overlap_grad_reduce=True - # +actor_rollout_ref.actor.megatron.override_ddp_config.overlap_param_gather=True - actor_rollout_ref.actor.megatron.use_distributed_optimizer=True - "++actor_rollout_ref.actor.megatron.override_transformer_config.pipeline_model_parallel_layout='${PIPELINE_MODEL_PARALLEL_LAYOUT}'" -) - -if [ "${RECOMPUTE_GRANULARITY}" = "selective" ]; then - RECOMPUTE_ARGS=( - ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=selective - "++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_modules=[mla_up_proj]" - ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=null - ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=null - ) -else - RECOMPUTE_ARGS=( - ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full - ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform - ++actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 - ) -fi - -# Context parallelism needs three extra transformer-config settings for DeepSeek-V4 on top of -# `context_parallel_size`. Each of them is enforced by Megatron-Core, so without them the run -# aborts at model build or in the first attention forward. They are appended only when CP > 1. -# -# cp_partition_mode=contiguous -# DSv4 attention requires every CP rank to own ONE consecutive interval of the packed THD -# buffer. Megatron-Core defaults to "zigzag" and raises -# "DSv4 Hybrid with CP requires cp_partition_mode='contiguous'." -# sequence_packing_scheduler=dp_balanced -# Megatron-Core asserts "DSv4 Hybrid with CP requires a sequence_packing_scheduler for THD -# inputs." Needs Transformer Engine >= 2.9. -# max_seqlen_per_dp_cp_rank -# Documented as max sequence length / cp_size; it drives how sub-samples are assigned to -# each DPxCP rank. -CP_ARGS=() -if [ "${ACTOR_CP}" -gt 1 ]; then - CP_ARGS=( - ++actor_rollout_ref.actor.megatron.override_transformer_config.cp_partition_mode=contiguous - ++actor_rollout_ref.actor.megatron.override_transformer_config.sequence_packing_scheduler=dp_balanced - ++actor_rollout_ref.actor.megatron.override_transformer_config.max_seqlen_per_dp_cp_rank=$(((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH) / ACTOR_CP)) - ) -fi - -# 截断重要性采样, 对 rollout/trainer 的 logprob 比值做裁剪。 -TIS_ARGS=() -if [ "${ROLLOUT_IS_MODE}" != "off" ]; then - TIS_ARGS=( - algorithm.rollout_correction.rollout_is=${ROLLOUT_IS_MODE} - algorithm.rollout_correction.rollout_is_threshold=${ROLLOUT_IS_THRESHOLD} - ) -fi - -ROLLOUT=( - actor_rollout_ref.rollout.name=vllm - actor_rollout_ref.rollout.tensor_model_parallel_size=1 - actor_rollout_ref.rollout.data_parallel_size=${ROLLOUT_EP} - actor_rollout_ref.rollout.expert_parallel_size=${ROLLOUT_EP} - actor_rollout_ref.rollout.enforce_eager=${ROLLOUT_ENFORCE_EAGER} - actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL} - actor_rollout_ref.rollout.n=${ROLLOUT_N} - actor_rollout_ref.rollout.val_kwargs.top_p=0.7 - actor_rollout_ref.rollout.val_kwargs.temperature=1.0 - actor_rollout_ref.rollout.val_kwargs.n=${ROLLOUT_N_VAL} - actor_rollout_ref.rollout.val_kwargs.do_sample=True - actor_rollout_ref.rollout.calculate_log_probs=True - actor_rollout_ref.rollout.enable_rollout_routing_replay=True - actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True - actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU} - actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 - actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} - actor_rollout_ref.rollout.max_model_len=${ROLLOUT_MAX_MODEL_LEN} - actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${ROLLOUT_UPDATE_WEIGHTS_BUCKET_MB} - +actor_rollout_ref.rollout.engine_kwargs.vllm.kv_cache_dtype=${ROLLOUT_KV_CACHE_DTYPE} - +actor_rollout_ref.rollout.quantization=fp8 -) - -REWARD=( - reward.reward_manager.name=dapo - +reward.reward_kwargs.overlong_buffer_cfg.enable=${OVERLONG_BUFFER_ENABLE} - +reward.reward_kwargs.overlong_buffer_cfg.len=${OVERLONG_BUFFER_LEN} - +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${OVERLONG_PENALTY_FACTOR} - +reward.reward_kwargs.overlong_buffer_cfg.log=False - +reward.reward_kwargs.max_resp_len=${MAX_RESPONSE_LENGTH} -) - -TRAINER=( - trainer.logger='["console","wandb"]' - trainer.project_name=${PROJECT_NAME} - trainer.experiment_name=${EXPERIMENT_NAME} - trainer.n_gpus_per_node=${NGPUS_PER_NODE} - trainer.nnodes=${NNODES} - trainer.save_freq=${SAVE_FREQ} - trainer.test_freq=${TEST_FREQ} - trainer.max_actor_ckpt_to_keep=2 - trainer.total_epochs=${TOTAL_EPOCHS} - trainer.total_training_steps=${TOTAL_TRAINING_STEPS} - trainer.resume_mode=auto - trainer.val_before_train=False - trainer.log_val_generations=100 - trainer.rollout_data_dir="${ROLLOUT_DATA_DIR}" - trainer.validation_data_dir="${VALIDATION_DATA_DIR}" - trainer.default_local_dir="${CKPTS_DIR}" -) - -EXTRA=( - actor_rollout_ref.nccl_timeout=3600 - model_engine=megatron -) - -# Explicitly disable QAT for both actor and rollout in this length-only recipe. -QAT_ARGS=( - ++actor_rollout_ref.actor.megatron.qat.enable=False - ++actor_rollout_ref.rollout.qat.enable=False -) - -########################### launch ########################### - -python3 -m verl.trainer.main_ppo \ - "${ALGORITHM[@]}" \ - "${DATA[@]}" \ - "${MODEL[@]}" \ - "${ACTOR[@]}" \ - "${RECOMPUTE_ARGS[@]}" \ - "${CP_ARGS[@]}" \ - "${TIS_ARGS[@]}" \ - "${ROLLOUT[@]}" \ - "${REWARD[@]}" \ - "${TRAINER[@]}" \ - "${EXTRA[@]}" \ - "${QAT_ARGS[@]}" \ - "$@" diff --git a/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh index 2ec004d952e..36152686f00 100755 --- a/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh +++ b/examples/grpo_trainer/run_deepseek_v4_flash_length_4096_qat_megatron.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash # DeepSeek-V4-Flash-0731: MXFP4 QAT, response 4096, 40 steps, validation every 10. -# Standalone MXFP4 QAT arm; the length-only baseline is a separate script. +# Standalone MXFP4 expert QAT recipe with TE FP8 training. # Plan A: prompt/response 1024/4096, overlong buffer 512 with penalty 1.0, # clip high 0.28 / dual clip 10.0, no thinking; expert MXFP4 QAT plus TE FP8 training. # Set MODEL_PATH, TRAIN_FILE and TEST_FILE in the launch environment. @@ -8,6 +8,41 @@ # For two segments, use TOTAL_TRAINING_STEPS=20, then 40 with the same checkpoint directory. # Requires DeepSeek-V4 support in Megatron-Bridge, Megatron-Core and vLLM. +# Megatron prerequisite: HDO with native FP32 parameters +# +# This DeepSeek-V4-Flash example enables `optimizer_cpu_offload=True` and retains some model +# parameters in FP32. For this combination, Megatron's `DistributedOptimizer` must detach FP32 +# model parameters before creating their optimizer shards. Otherwise, reconstructing +# `HybridDeviceOptimizer` can fail during optimizer initialization with: +# +# ```text +# ValueError: can't optimize a non-leaf Tensor +# ``` +# +# The required fix is in `megatron/core/optimizer/distrib_optimizer.py`, inside +# `DistributedOptimizer._build_model_and_main_param_groups`, in the FP32-parameter branch: +# +# ```diff +# - 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] +# ``` +# +# Use a Megatron version containing this fix, or backport it to the selected dependency checkout. +# The upstream fix is [NVIDIA/Megatron-LM#6982](https://github.com/NVIDIA/Megatron-LM/pull/6982), +# commit +# [`d5ff7ea72cffe1eb1daedaa2c7b2858694ab0d3c`](https://github.com/NVIDIA/Megatron-LM/commit/d5ff7ea72cffe1eb1daedaa2c7b2858694ab0d3c). +# +# **Validation provenance:** the completed 40-step MXFP4 QAT experiment used Megatron +# `1ff25ca7e339fe521165da7f4373d9f52e7af436` with an equivalent local runtime patch, +# `_patch_nonleaf_fp32_shards`. Both the initial 20-step segment and the resumed segment to step +# 40 logged that the patch was applied. Validation ran at steps 10/20/30/40, with checkpoints +# saved at steps 20/40 and optimizer state restored between segments. +# +# The runtime patch helper is not included in the published verl changes. Reproducing the +# experiment with Megatron `1ff25ca7e` therefore requires the dependency fix above in addition to +# the verl changes. This optimizer-initialization prerequisite is separate from the MXFP4 +# checkpoint serialization and HDO checkpoint-resume fixes in verl. + set -euo pipefail export CUDA_DEVICE_MAX_CONNECTIONS=1 From 4e811b54ad00df25cfc755ad2966dd7f16eb4dbf Mon Sep 17 00:00:00 2001 From: Sophia Yang Date: Wed, 16 Sep 2026 07:57:26 -0700 Subject: [PATCH 6/6] [example] fix QAT recipe naming and generated config Rename the original DeepSeek-V4-Flash QAT example to use the required Megatron backend suffix, preserving its contents. Refresh the generated Megatron reference config with the existing bypass_te_fp8_assert defaults. Co-authored-by: Codex --- ...h_mcore_fp4qat.sh => run_deepseek_v4_flash_qat_megatron.sh} | 0 verl/trainer/config/_generated_ppo_megatron_trainer.yaml | 3 +++ 2 files changed, 3 insertions(+) rename examples/grpo_trainer/{run_deepseek_v4_flash_mcore_fp4qat.sh => run_deepseek_v4_flash_qat_megatron.sh} (100%) diff --git a/examples/grpo_trainer/run_deepseek_v4_flash_mcore_fp4qat.sh b/examples/grpo_trainer/run_deepseek_v4_flash_qat_megatron.sh similarity index 100% rename from examples/grpo_trainer/run_deepseek_v4_flash_mcore_fp4qat.sh rename to examples/grpo_trainer/run_deepseek_v4_flash_qat_megatron.sh diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index bd912745aea..d94a4f75584 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -97,6 +97,7 @@ actor_rollout_ref: - '*mlp.gate' activation_observer: null quantization_config_path: null + bypass_te_fp8_assert: false _target_: verl.workers.config.McoreActorConfig rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1} strategy: megatron @@ -300,6 +301,7 @@ actor_rollout_ref: - '*mlp.gate' activation_observer: null quantization_config_path: null + bypass_te_fp8_assert: false _target_: verl.workers.config.McoreActorConfig entropy_from_logits_with_chunking: ${oc.select:actor_rollout_ref.actor.entropy_from_logits_with_chunking,False} entropy_from_logits_chunk_size: ${oc.select:actor_rollout_ref.actor.entropy_from_logits_chunk_size,2048} @@ -639,6 +641,7 @@ critic: - '*mlp.gate' activation_observer: null quantization_config_path: null + bypass_te_fp8_assert: false _target_: verl.workers.config.McoreCriticConfig rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1} strategy: megatron