From 489a51082bb39499e7dd5b9fec6b08d073e773d2 Mon Sep 17 00:00:00 2001 From: mouad-hpc Date: Wed, 17 Jun 2026 11:52:54 -0700 Subject: [PATCH 01/14] feat: block-FP8 LoRA colocate support for Qwen3.5-35B-A3B --- ...n-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh | 172 ++++++++++++++++++ .../update_weight_from_tensor.py | 38 ++-- 2 files changed, 192 insertions(+), 18 deletions(-) create mode 100755 examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh diff --git a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh new file mode 100755 index 00000000000..d3947a24b35 --- /dev/null +++ b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh @@ -0,0 +1,172 @@ +#!/bin/bash + +# Qwen3.5-35B-A3B GRPO with block-wise FP8 (e4m3) base + o_proj LoRA, colocate. + +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex +export PYTHONBUFFERED=16 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi + +source "/root/miles/scripts/models/qwen3.5-35B-A3B.sh" + +CKPT_ARGS=( + # Bridge-load the base from HF; --ref-load (torch_dist) crashes on this + # hybrid-GDN model's _extra_state. + --hf-checkpoint /root/Qwen3.5-35B-A3B-FP8 +) + +LORA_ARGS=( + --lora-rank 32 + --lora-alpha 32 + --lora-dropout 0.0 + # o_proj is the one unfused/ungated attention projection; q/k/v/MoE targets + # need the gated qkv-LoRA buffer fix upstream in SGLang. + --lora-type lora + --target-modules "o_proj" + --sglang-lora-backend triton + --megatron-to-hf-mode bridge +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 4096 + --rollout-temperature 1 + --global-batch-size 256 + --balance-data +) + +EVAL_ARGS=( + --skip-eval-before-train + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 1 + --eval-max-response-len 8000 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # GDN rejects packed sequences; bshd pads per-sequence (needs static batches). + --qkv-format bshd + --micro-batch-size 1 + + --moe-enable-deepep + --moe-token-dispatcher-type flex + + --transformer-impl transformer_engine + --bf16 + --fp8-format e4m3 + --fp8-recipe blockwise +) + +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( +) + +SGLANG_ARGS=( + # Block-FP8 [128,128] needs every sharded dim a multiple of 128, so the + # shared-expert MLP caps world-TP at 4 => 2 engines x 4 GPUs. + --rollout-num-gpus-per-engine 4 + --sglang-mem-fraction-static 0.4 + # ep=1 avoids the -1 non-local-expert sentinels that trip the MoE-LoRA align kernel. + --sglang-ep-size 1 + --sglang-disable-cuda-graph + --sglang-dtype bfloat16 + --sglang-max-running-requests 512 + --sglang-moe-runner-backend triton + # Preserve the FP8 base across the colocate torch_memory_saver release/resume: + # resume() discards weight content, and MILES_SKIP_BASE_SYNC keeps the base from + # being re-synced, so without this the base is corrupted (gibberish rollouts). + --sglang-enable-weights-cpu-backup +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --update-weight-buffer-size 536870912 +) + +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", + \"NVTE_FP8_BLOCK_SCALING_FP32_SCALES\": \"1\", + \"NCCL_TIMEOUT_MS\": \"36000000\", + \"MILES_SKIP_BASE_SYNC\": \"1\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --colocate \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${LORA_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${SGLANG_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 7c82d0f122c..887922e24d8 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -1,4 +1,5 @@ import logging +import os from argparse import Namespace from collections.abc import Callable, Mapping, Sequence from contextlib import nullcontext @@ -198,8 +199,13 @@ def update_weights(self) -> None: megatron_local_weights = self.weights_getter() - # For LoRA+distributed: base weights are frozen, skip after first round. - if not (self.is_lora and self.use_distribute and self._lora_base_synced): + # LoRA base is frozen and SGLang serves it from the FP8 checkpoint, so the + # per-round base re-export is redundant. MILES_SKIP_BASE_SYNC extends the + # existing distributed-LoRA base skip to colocate; only adapters sync. + skip_base_sync = (self.is_lora and self.use_distribute and self._lora_base_synced) or ( + (self.is_lora or self.is_multi_lora) and os.environ.get("MILES_SKIP_BASE_SYNC") == "1" + ) + if not skip_base_sync: base_ctx = nullcontext() if self.is_multi_lora: # For multi_lora, hide the multi-adapter layer entirely so it doesn't @@ -242,13 +248,10 @@ def update_weights(self) -> None: dist.barrier(group=get_gloo_group()) if rank == 0: - # `post_process_quantization` is related to the `process_weights_after_loading` - # in the sglang rollout side, which should always be invoked after weight - # updating. post_process_weights( rollout_engines=self.rollout_engines, restore_weights_before_load=False, - post_process_quantization=True, + post_process_quantization=not skip_base_sync, ) ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) @@ -418,19 +421,18 @@ def _send_to_colocated_engine( if lora_loaded: ray.get(ipc_engine.unload_lora_adapter.remote(lora_name=lora_name)) - # (Yusheng) to-do-1: update lora weights from tensors should support multiple dtypes (bf16, fp8, fp16, fp32) - # currently, we only support 1 type. If there are multiple dtypes, we need to serialize the tensors for each dtype. - # Thus, we need to apply the same way as `ipc_engine.update_weights_from_tensor` in future - # (Yusheng) to-do-2: need to add ci test acc here - now it will pass but fail to update lora weights - - refs.append( - ipc_engine.load_lora_adapter_from_tensors.remote( - lora_name=lora_name, - config_dict=lora_config, - serialized_tensors=serialized_named_tensors[0][0], - load_format="flattened_bucket", + # Loop the dtype index (rank-replicated, so rank stays [0]); SGLang merges + # successive loads under the same lora_name into one adapter. + num_dtypes = len(serialized_named_tensors[0]) + for i in range(num_dtypes): + refs.append( + ipc_engine.load_lora_adapter_from_tensors.remote( + lora_name=lora_name, + config_dict=lora_config, + serialized_tensors=serialized_named_tensors[0][i], + load_format="flattened_bucket", + ) ) - ) else: num_dtypes = len(serialized_named_tensors[0]) From 82f1877c6033e32a3dd2e30b9971ad66e0dbe02e Mon Sep 17 00:00:00 2001 From: mouad-hpc Date: Mon, 22 Jun 2026 17:00:25 -0700 Subject: [PATCH 02/14] feat: gated canonical q/k/v/o LoRA for FP8 Qwen3.5 (gate-aware q=8192, bridge ModuleDict + interleave fixes) --- ...n-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh | 70 +++++++---- miles/backends/megatron_utils/lora_utils.py | 3 + .../update_weight_from_tensor.py | 13 +- miles_plugins/megatron_bridge/__init__.py | 3 + .../megatron_bridge/gated_canonical_lora.py | 113 ++++++++++++++++++ 5 files changed, 176 insertions(+), 26 deletions(-) create mode 100644 miles_plugins/megatron_bridge/gated_canonical_lora.py diff --git a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh index d3947a24b35..0a22c9d35ce 100755 --- a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh +++ b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh @@ -1,7 +1,17 @@ #!/bin/bash -# Qwen3.5-35B-A3B GRPO with block-wise FP8 (e4m3) base + o_proj LoRA, colocate. - +# Qwen3.5-35B-A3B MoE LoRA + block-wise FP8 training (Hopper / Blackwell). +# +# Combines: +# - LoRA on MoE expert projections (gate_proj, up_proj, down_proj) +# - SGLang triton LoRA backend (required for MoE LoRA) +# - Megatron-Bridge HF conversion (required for LoRA path) +# - Block-wise FP8 e4m3 forward, BF16 backward + master weights +# - --use-tis for MoE numerical drift compensation +# +# See docs/superpowers/plans/2026-06-08-fp8-moe-lora-02-fp8-moe-lora-bringup.md. + +# for rerun the task pkill -9 sglang sleep 3 ray stop --force @@ -12,6 +22,8 @@ pkill -9 ray pkill -9 python set -ex + +# will prevent ray from buffering stdout/stderr export PYTHONBUFFERED=16 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) @@ -20,25 +32,30 @@ if [ "$NVLINK_COUNT" -gt 0 ]; then else HAS_NVLINK=0 fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" source "/root/miles/scripts/models/qwen3.5-35B-A3B.sh" CKPT_ARGS=( - # Bridge-load the base from HF; --ref-load (torch_dist) crashes on this - # hybrid-GDN model's _extra_state. + # Bridge mode loads the base model from --hf-checkpoint via Megatron-Bridge. + # No --ref-load: a torch_dist load routes through Megatron dist_checkpointing, + # which crashes on this hybrid-GDN model's _extra_state + # (_replace_sharded_keys_with_state_dict_keys: "BytesIO has no len()"), + # regardless of which image built the torch_dist. The canonical MoE-LoRA + # bridge recipe (run-gpt-oss-20B-megatron-moe-lora.sh) loads from HF instead. --hf-checkpoint /root/Qwen3.5-35B-A3B-FP8 ) LORA_ARGS=( - --lora-rank 32 - --lora-alpha 32 - --lora-dropout 0.0 - # o_proj is the one unfused/ungated attention projection; q/k/v/MoE targets - # need the gated qkv-LoRA buffer fix upstream in SGLang. - --lora-type lora - --target-modules "o_proj" + --lora-rank 32 # LoRA rank + --lora-alpha 32 # LoRA alpha (= rank for RL) + --lora-dropout 0.0 # 0 for RL + # canonical_lora exports separate q/k/v so SGLang applies them unfused; + # gated_canonical_lora sizes the gated q adapter to 8192 (query+gate). + --lora-type canonical_lora + --target-modules "q_proj,k_proj,v_proj,o_proj" --sglang-lora-backend triton - --megatron-to-hf-mode bridge + --megatron-to-hf-mode bridge # required for LoRA path ) ROLLOUT_ARGS=( @@ -51,8 +68,10 @@ ROLLOUT_ARGS=( --num-rollout 3000 --rollout-batch-size 32 --n-samples-per-prompt 8 + # 4096 avoids the fp32-logits train-step OOM at 8192 on colocated H200s. --rollout-max-response-len 4096 --rollout-temperature 1 + --global-batch-size 256 --balance-data ) @@ -78,13 +97,15 @@ PERF_ARGS=( --recompute-method uniform --recompute-num-layers 1 - # GDN rejects packed sequences; bshd pads per-sequence (needs static batches). + # GDN rejects packed sequences; bshd pads per-sequence (needs static micro batches). --qkv-format bshd --micro-batch-size 1 + # use deepep for megatron MoE --moe-enable-deepep --moe-token-dispatcher-type flex + # block-wise FP8 --transformer-impl transformer_engine --bf16 --fp8-format e4m3 @@ -98,7 +119,7 @@ GRPO_ARGS=( --entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 - --use-tis + --use-tis # MoE precision-drift compensation ) OPTIMIZER_ARGS=( @@ -111,22 +132,26 @@ OPTIMIZER_ARGS=( ) WANDB_ARGS=( + # --use-wandb + # --wandb-project miles-fp8-moe-lora + # --wandb-group qwen3.5-35B-A3B-fp8-moe-lora + # --wandb-key ${WANDB_KEY} ) SGLANG_ARGS=( - # Block-FP8 [128,128] needs every sharded dim a multiple of 128, so the - # shared-expert MLP caps world-TP at 4 => 2 engines x 4 GPUs. + # Block-FP8 needs every sharded dim a multiple of 128. shared_expert=512 and + # moe_ffn=512 cap the per-engine TP at 4, so use 2 engines x 4 GPUs. ep=1 + # avoids the MoE-LoRA align-kernel IMA on EP's -1 expert sentinels. --rollout-num-gpus-per-engine 4 --sglang-mem-fraction-static 0.4 - # ep=1 avoids the -1 non-local-expert sentinels that trip the MoE-LoRA align kernel. --sglang-ep-size 1 + # Hybrid-GDN cuda-graph capture deadlocks under colocate; run eager. --sglang-disable-cuda-graph --sglang-dtype bfloat16 + + --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256) --sglang-max-running-requests 512 --sglang-moe-runner-backend triton - # Preserve the FP8 base across the colocate torch_memory_saver release/resume: - # resume() discards weight content, and MILES_SKIP_BASE_SYNC keeps the base from - # being re-synced, so without this the base is corrupted (gibberish rollouts). --sglang-enable-weights-cpu-backup ) @@ -136,12 +161,15 @@ MISC_ARGS=( --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 --attention-backend flash - --update-weight-buffer-size 536870912 + --update-weight-buffer-size 536870912 # 512MB ) +# launch the master node of ray in container export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +# NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 forces fp32 scales in fp8 training, +# matching what sglang serves on the rollout side. RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"/root/Megatron-LM/\", diff --git a/miles/backends/megatron_utils/lora_utils.py b/miles/backends/megatron_utils/lora_utils.py index 6820202e861..cc0e92ef3cf 100644 --- a/miles/backends/megatron_utils/lora_utils.py +++ b/miles/backends/megatron_utils/lora_utils.py @@ -249,6 +249,9 @@ def create_lora_instance(args: Namespace): Returns: A LoRA/CanonicalLoRA dataclass instance ready to be applied to a model. """ + # Install gate-aware CanonicalLoRA.transform before the LoRA is built. + import miles_plugins.megatron_bridge # noqa: F401 + from megatron.bridge.peft.canonical_lora import CanonicalLoRA from megatron.bridge.peft.lora import LoRA diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 887922e24d8..c749caf68ed 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -199,9 +199,9 @@ def update_weights(self) -> None: megatron_local_weights = self.weights_getter() - # LoRA base is frozen and SGLang serves it from the FP8 checkpoint, so the - # per-round base re-export is redundant. MILES_SKIP_BASE_SYNC extends the - # existing distributed-LoRA base skip to colocate; only adapters sync. + # LoRA base is frozen and SGLang serves it from the FP8 checkpoint; re-syncing + # it every round is redundant and re-quantization corrupts the block-FP8 weights. + # MILES_SKIP_BASE_SYNC extends the distributed skip to colocate; only adapters sync. skip_base_sync = (self.is_lora and self.use_distribute and self._lora_base_synced) or ( (self.is_lora or self.is_multi_lora) and os.environ.get("MILES_SKIP_BASE_SYNC") == "1" ) @@ -248,6 +248,10 @@ def update_weights(self) -> None: dist.barrier(group=get_gloo_group()) if rank == 0: + # process_weights_after_loading is not idempotent for block-FP8, so only + # re-quantize when the base was actually re-synced. The call itself is kept + # unconditionally: it drains the engine queue so the async adapter load + # finishes before the trainer frees its CUDA-IPC source tensors. post_process_weights( rollout_engines=self.rollout_engines, restore_weights_before_load=False, @@ -421,8 +425,7 @@ def _send_to_colocated_engine( if lora_loaded: ray.get(ipc_engine.unload_lora_adapter.remote(lora_name=lora_name)) - # Loop the dtype index (rank-replicated, so rank stays [0]); SGLang merges - # successive loads under the same lora_name into one adapter. + # One load per dtype bucket; SGLang merges them server-side into one adapter. num_dtypes = len(serialized_named_tensors[0]) for i in range(num_dtypes): refs.append( diff --git a/miles_plugins/megatron_bridge/__init__.py b/miles_plugins/megatron_bridge/__init__.py index e69de29bb2d..f910cb8fb0e 100644 --- a/miles_plugins/megatron_bridge/__init__.py +++ b/miles_plugins/megatron_bridge/__init__.py @@ -0,0 +1,3 @@ +from . import gated_canonical_lora + +gated_canonical_lora.install() diff --git a/miles_plugins/megatron_bridge/gated_canonical_lora.py b/miles_plugins/megatron_bridge/gated_canonical_lora.py new file mode 100644 index 00000000000..bf351921c43 --- /dev/null +++ b/miles_plugins/megatron_bridge/gated_canonical_lora.py @@ -0,0 +1,113 @@ +"""Gate-aware q sizing for CanonicalLoRA. + +Megatron-Bridge's ``CanonicalLoRA.transform`` sizes the q adapter as +``kv_channels * num_attention_heads``, correct only for ungated attention. +Qwen3.5 uses gated attention, so ``q_proj`` outputs query + output gate +(``2 * num_attention_heads * head_dim``); this mirrors the bridge transform but +doubles ``q_out_features`` when ``attention_output_gate`` is set. The formula +reduces to the original for ungated models. +""" + +from __future__ import annotations + +import logging + +import megatron.bridge.peft.canonical_lora as cl +import torch +from torch import nn + +logger = logging.getLogger(__name__) + + +def interleave_qkv_gated(self, query, key, value): + """Gate-aware replacement for ``LoRALinearSplitQKV._interleave_qkv``. + + The base implementation sizes q heads from ``num_attention_heads`` and cannot + place the gated query (2x heads) into Megatron's per-group qkv layout. + """ + config = self.to_wrap.config + head_dim = config.kv_channels + num_kv = config.num_query_groups + q_per_group = config.num_attention_heads // num_kv + gate = 2 if getattr(config, "attention_output_gate", False) else 1 + lead = query.shape[:-1] + q = ( + query.reshape(*lead, num_kv, q_per_group, gate, head_dim) + .transpose(-3, -2) + .reshape(*lead, num_kv, gate * q_per_group * head_dim) + ) + k = key.reshape(*lead, num_kv, head_dim) + v = value.reshape(*lead, num_kv, head_dim) + return torch.cat([q, k, v], dim=-1).reshape(*lead, -1) + + +def patched_transform(self, m, name=None, prefix=None): + if isinstance(m, (cl.LinearAdapter, cl.LoRALinear, cl.LoRALinearSplitQKV, cl.LoRALinearSplitFC1UpGate)): + return m + + ans = self.match(m, name, prefix) + if ans is None: + return m + match, full_name = ans + + if isinstance(m, nn.Linear): + return cl.LinearAdapter( + m, dim=self.dim, alpha=self.alpha, dropout=self.dropout, lora_A_init_method=self.lora_A_init_method + ) + + is_expert = cl.is_expert_linear(full_name) + attrs = cl.get_adapter_attributes_from_linear(m, is_expert=is_expert) + + adapter_kwargs = dict( + dim=self.dim, + base_linear_name=full_name, + activation="identity", + norm_type=None, + column_init_method=self.lora_A_init_method, + row_init_method=self.lora_B_init_method, + gather_output=False, + input_is_parallel=attrs.input_is_parallel, + dropout=self.dropout, + dropout_position=self.dropout_position, + model_parallel_config=getattr(m, "config", None), + alpha=self.alpha, + is_expert=is_expert, + disable_tensor_parallel_comm=attrs.disable_tensor_parallel_comm, + disable_sequence_parallel_comm=attrs.disable_sequence_parallel_comm, + base_linear_is_parallel=attrs.base_linear_is_parallel, + ) + + if name in ["linear_proj", "linear_fc2"]: + adapter = cl.ParallelLinearAdapter(attrs.in_features, attrs.out_features, **adapter_kwargs) + return cl.LoRALinear(m, adapter) + + canonical_submodules = self.canonical_mapping[match] + if name == "linear_qkv": + adapter_q = adapter_k = adapter_v = None + kv_out_features = m.config.kv_channels * m.config.num_query_groups + q_out_features = m.config.kv_channels * m.config.num_attention_heads + if getattr(m.config, "attention_output_gate", False): + q_out_features *= 2 + if "linear_q" in canonical_submodules: + adapter_q = cl.ParallelLinearAdapter(attrs.in_features, q_out_features, **adapter_kwargs) + if "linear_k" in canonical_submodules: + adapter_k = cl.ParallelLinearAdapter(attrs.in_features, kv_out_features, **adapter_kwargs) + if "linear_v" in canonical_submodules: + adapter_v = cl.ParallelLinearAdapter(attrs.in_features, kv_out_features, **adapter_kwargs) + return cl.LoRALinearSplitQKV(m, cl.ModuleDict({"adapter_q": adapter_q, "adapter_k": adapter_k, "adapter_v": adapter_v})) + + if name == "linear_fc1": + adapter_up = adapter_gate = None + if "linear_fc1_up" in canonical_submodules: + adapter_up = cl.ParallelLinearAdapter(attrs.in_features, attrs.out_features // 2, **adapter_kwargs) + if "linear_fc1_gate" in canonical_submodules: + adapter_gate = cl.ParallelLinearAdapter(attrs.in_features, attrs.out_features // 2, **adapter_kwargs) + return cl.LoRALinearSplitFC1UpGate(m, cl.ModuleDict({"adapter_up": adapter_up, "adapter_gate": adapter_gate})) + + return m + + +def install() -> None: + cl.CanonicalLoRA.transform = patched_transform + cl.LoRALinearSplitQKV._interleave_qkv = interleave_qkv_gated + logger.info("Installed gate-aware q sizing for CanonicalLoRA.transform") From 075de9f73f18c7f509076aa099b700dc00b0e256 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Mon, 13 Jul 2026 15:41:59 -0700 Subject: [PATCH 03/14] forward fp8 GEMM-autocast args to the bridge provider --- miles/backends/megatron_utils/model_provider.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 45adc86792f..024e18ea5b6 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -110,6 +110,10 @@ def wrapped_model_provider( provider.moe_router_bias_update_rate = args.moe_router_bias_update_rate if getattr(args, "moe_aux_loss_coeff", None) is not None: provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff + # AutoBridge derives the provider from the HF config only, so Megatron's fp8 + # GEMM-autocast flags are dead here unless forwarded. + provider.fp8 = args.fp8 + provider.fp8_recipe = args.fp8_recipe provider.finalize() def wrapped_bridge_provider( From f2b33d73536e9da4265c5457a88e02865e8876d2 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Mon, 13 Jul 2026 10:42:58 -0700 Subject: [PATCH 04/14] fp8: post-checkpoint block-fp8 storage of the frozen LoRA base --- .../megatron_utils/fp8_frozen_base.py | 130 ++++++++++++++++++ miles/backends/megatron_utils/model.py | 4 + miles/utils/arguments.py | 9 ++ 3 files changed, 143 insertions(+) create mode 100644 miles/backends/megatron_utils/fp8_frozen_base.py diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py new file mode 100644 index 00000000000..f35764547f2 --- /dev/null +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -0,0 +1,130 @@ +import logging + +import torch +import torch.distributed as dist + +from miles.backends.megatron_utils.lora_utils import _is_adapter_param_name +from miles.utils.fp8_kernel import blockwise_cast_to_fp8_triton + +from tools.fp8_cast_bf16 import weight_dequant + +logger = logging.getLogger(__name__) + +BLOCK = 128 + +# Frozen base linear weights to store as block-fp8 (mirrors the export-side +# allowlist in megatron_to_hf/processors/quantizer_fp8.py, plus GDN in/out proj). +LINEAR_SUBSTR = ( + "self_attention.linear_qkv", + "self_attention.linear_proj", + "self_attention.linear_q_proj", + "self_attention.linear_q_down_proj", + "self_attention.linear_q_up_proj", + "self_attention.linear_kv_down_proj", + "self_attention.linear_kv_up_proj", + "mlp.linear_fc1", + "mlp.linear_fc2", + "mlp.shared_experts.linear_fc1", + "mlp.shared_experts.linear_fc2", + "mlp.experts.linear_fc1", + "mlp.experts.linear_fc2", + "linear_attn.in_proj_a", + "linear_attn.in_proj_b", + "linear_attn.in_proj_qkv", + "linear_attn.in_proj_z", + "linear_attn.out_proj", +) + + +def rank0() -> bool: + return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0 + + +def is_base_linear_weight(name: str) -> bool: + leaf = name.rsplit(".", 1)[-1] + if not (leaf == "weight" or (leaf.startswith("weight") and leaf[len("weight") :].isdigit())): + return False + if _is_adapter_param_name(name): + return False + return any(s in name for s in LINEAR_SUBSTR) + + +def should_quantize(name: str, param: torch.nn.Parameter) -> bool: + return param is not None and param.dim() == 2 and not param.requires_grad and is_base_linear_weight(name) + + +def family(name: str) -> str: + if "mlp.experts." in name: + return "moe_experts" + if "shared_experts" in name: + return "shared_experts" + if "linear_attn" in name: + return "gdn" + if "self_attention" in name: + return "attention" + if "mlp.linear_fc" in name: + return "dense_mlp" + return "other" + + +def install_fp8_hooks(module: torch.nn.Module) -> None: + def pre(mod, inputs): + for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): + q = getattr(mod, f"fp8q_{leaf}").contiguous() + s = getattr(mod, f"fp8s_{leaf}").contiguous() + param.data = weight_dequant(q, s, BLOCK).to(dtype).reshape(shape) + + def post(mod, inputs, output): + # Free the transient bf16. Correct because the base is frozen (no wgrad) + # and TE fp8 / activation-recompute keeps whatever backward needs; for a + # plain layer the autograd graph still holds its own saved copy. + for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): + param.data = torch.empty(0, dtype=dtype, device=param.data.device) + return output + + module.register_forward_pre_hook(pre) + module.register_forward_hook(post) + + +def quantize_frozen_base_to_fp8(model_chunks, args) -> None: + freed_bytes = 0 + counts: dict[str, int] = {} + roundtrip_relerr = None + + for chunk in model_chunks: + for mod_name, module in chunk.named_modules(): + entries: dict = {} + for leaf, param in list(module.named_parameters(recurse=False)): + full = f"{mod_name}.{leaf}" if mod_name else leaf + if not should_quantize(full, param): + continue + + w = param.data.contiguous() + q, s = blockwise_cast_to_fp8_triton(w, [BLOCK, BLOCK]) + + if roundtrip_relerr is None: + recon = weight_dequant(q.contiguous(), s.contiguous(), BLOCK).to(w.dtype) + denom = w.abs().amax().clamp(min=1e-6) + roundtrip_relerr = ((recon - w).abs().amax() / denom).item() + + module.register_buffer(f"fp8q_{leaf}", q, persistent=False) + module.register_buffer(f"fp8s_{leaf}", s, persistent=False) + entries[leaf] = (param, tuple(param.shape), param.dtype) + + freed_bytes += w.numel() * w.element_size() + freed_bytes -= q.numel() * q.element_size() + s.numel() * s.element_size() + counts[family(full)] = counts.get(family(full), 0) + 1 + param.data = torch.empty(0, dtype=param.dtype, device=param.device) + + if entries: + module.fp8_frozen_entries = entries + install_fp8_hooks(module) + + if rank0(): + logger.info( + "fp8 frozen base: quantized %d tensors %s, freed ~%.2f GB/rank, roundtrip_relerr=%s", + sum(counts.values()), + counts, + freed_bytes / (1024**3), + f"{roundtrip_relerr:.2e}" if roundtrip_relerr is not None else "n/a", + ) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 77376a5af63..140e8c9605f 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -816,6 +816,10 @@ def initialize_model_and_optimizer( checkpointing_context={}, skip_load_to_model_and_opt=False, ) + if getattr(args, "fp8_frozen_base_store", False): + from miles.backends.megatron_utils.fp8_frozen_base import quantize_frozen_base_to_fp8 + + quantize_frozen_base_to_fp8(model, args) check_peak_gpu_memory_after_load(args) clear_memory() diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index a842a21d099..2dd2e71bc74 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -211,6 +211,15 @@ def add_train_arguments(parser): parser.add_argument( "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) + parser.add_argument( + "--fp8-frozen-base-store", + action="store_true", + default=False, + help=( + "Post-checkpoint: store the frozen LoRA base linear weights as block-fp8 " + "(dequant per layer in the forward) to halve resident base weight memory." + ), + ) parser.add_argument( "--allgather-cp", action="store_true", From f1b499e568275cf109aeaa885cbdd5544307fba2 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Mon, 13 Jul 2026 12:05:42 -0700 Subject: [PATCH 05/14] fp8 frozen base: free the transient bf16 at offload, not after forward --- miles/backends/megatron_utils/actor.py | 12 +++++++-- .../megatron_utils/fp8_frozen_base.py | 26 ++++++++++++------- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 0e3490aca9a..3e77c56d42a 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -119,8 +119,8 @@ def init( initialize_multi_lora_model_and_optimizer(args, role) ) else: - (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = ( - initialize_model_and_optimizer(args, role) + (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer( + args, role ) parallel_state = get_parallel_state() @@ -206,6 +206,10 @@ def sleep(self) -> None: if not self.args.offload_train: return + if getattr(self.args, "fp8_frozen_base_store", False): + from miles.backends.megatron_utils.fp8_frozen_base import free_frozen_base + + free_frozen_base(self.model) clear_memory(clear_host_memory=True) print_memory("before offload model") destroy_process_groups() @@ -485,6 +489,7 @@ def load_pending_adapters(self) -> int: from miles.ray.multi_lora_controller import get_multi_lora_controller from miles.utils.adapter_config import AdapterState + configs = ray.get(get_multi_lora_controller().adapter_configs.remote()) if not any(c.state == AdapterState.PENDING for c in configs.values()): return 0 @@ -493,6 +498,7 @@ def load_pending_adapters(self) -> int: # self.wake_up() from .update_weight.multi_lora_sync import load_pending_adapters + n = load_pending_adapters(self.args, self.model, self.optimizer) if n > 0: # Re-snapshot: init_adapter_slot + load_adapter mutated the model, @@ -511,12 +517,14 @@ def unload_drained_adapters(self, rollout_id: int) -> int: return 0 from miles.ray.multi_lora_controller import get_multi_lora_controller from miles.utils.adapter_config import AdapterState + configs = ray.get(get_multi_lora_controller().adapter_configs.remote()) if not any(c.state == AdapterState.DRAINED for c in configs.values()): return 0 # if self.args.offload_train: # self.wake_up() from .update_weight.multi_lora_sync import unload_drained_adapters + n = unload_drained_adapters(self.args, self.model, self.optimizer, rollout_id) if n > 0: self.weights_backuper.backup("actor") diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index f35764547f2..1ec35e36a5d 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -68,22 +68,30 @@ def family(name: str) -> str: def install_fp8_hooks(module: torch.nn.Module) -> None: + # Materialize the bf16 weight before every forward. It is NOT freed after the + # forward: TE grouped-expert backward reads self.weight directly, and under + # activation recompute the recomputed forward + its backward straddle a free. + # The transient bf16 is released together at offload time (free_frozen_base). def pre(mod, inputs): for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): q = getattr(mod, f"fp8q_{leaf}").contiguous() s = getattr(mod, f"fp8s_{leaf}").contiguous() param.data = weight_dequant(q, s, BLOCK).to(dtype).reshape(shape) - def post(mod, inputs, output): - # Free the transient bf16. Correct because the base is frozen (no wgrad) - # and TE fp8 / activation-recompute keeps whatever backward needs; for a - # plain layer the autograd graph still holds its own saved copy. - for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): - param.data = torch.empty(0, dtype=dtype, device=param.data.device) - return output - module.register_forward_pre_hook(pre) - module.register_forward_hook(post) + + +def free_frozen_base(model_chunks) -> None: + """Drop the transient bf16 base weights so the offloaded/resident footprint is + the fp8 buffers. Call after the training step's backward (e.g. before offload). + The forward pre-hook re-materializes them on the next step.""" + for chunk in model_chunks: + for module in chunk.modules(): + entries = getattr(module, "fp8_frozen_entries", None) + if not entries: + continue + for leaf, (param, shape, dtype) in entries.items(): + param.data = torch.empty(0, dtype=dtype, device=param.data.device) def quantize_frozen_base_to_fp8(model_chunks, args) -> None: From aff250bc9dd9ec986385e0caf5f9bd8d5133f581 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Mon, 13 Jul 2026 14:32:44 -0700 Subject: [PATCH 06/14] exclude fp8 frozen base from weights backuper; cover bridge GDN naming --- miles/backends/megatron_utils/actor.py | 18 +++++++++++++++--- .../backends/megatron_utils/fp8_frozen_base.py | 14 +++++++++++++- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 3e77c56d42a..681a8ecead9 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -141,13 +141,25 @@ def init( start_rollout_id = loaded_rollout_id + 1 - self.weights_backuper = TensorBackuper.create( - source_getter=lambda: named_params_and_buffers( + def backup_source(): + source = named_params_and_buffers( self.args, self.model, convert_to_global_name=args.megatron_to_hf_mode == "raw", translate_gpu_to_cpu=not self.args.enable_weights_backuper, - ), + ) + if not self.args.fp8_frozen_base_store: + return source + from miles.backends.megatron_utils.fp8_frozen_base import frozen_fp8_param_ids + + # The frozen base lives as fp8 buffers; its bf16 param.data is transient + # (materialized by the forward pre-hook, freed at offload), so it must not + # be backed up or restored. + skip = frozen_fp8_param_ids(self.model) + return ((name, tensor) for name, tensor in source if id(tensor) not in skip) + + self.weights_backuper = TensorBackuper.create( + source_getter=backup_source, single_tag=None if args.enable_weights_backuper else "actor", ) self._active_model_tag: str | None = "actor" diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index 1ec35e36a5d..6346dd1e07c 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -33,6 +33,9 @@ "linear_attn.in_proj_qkv", "linear_attn.in_proj_z", "linear_attn.out_proj", + # bridge-built GDN layers keep the GatedDeltaNet module at self_attention + "self_attention.in_proj", + "self_attention.out_proj", ) @@ -58,7 +61,7 @@ def family(name: str) -> str: return "moe_experts" if "shared_experts" in name: return "shared_experts" - if "linear_attn" in name: + if "linear_attn" in name or "self_attention.in_proj" in name or "self_attention.out_proj" in name: return "gdn" if "self_attention" in name: return "attention" @@ -81,6 +84,15 @@ def pre(mod, inputs): module.register_forward_pre_hook(pre) +def frozen_fp8_param_ids(model_chunks) -> set[int]: + return { + id(param) + for chunk in model_chunks + for module in chunk.modules() + for param, _, _ in getattr(module, "fp8_frozen_entries", {}).values() + } + + def free_frozen_base(model_chunks) -> None: """Drop the transient bf16 base weights so the offloaded/resident footprint is the fp8 buffers. Call after the training step's backward (e.g. before offload). From 77890c11217bccf63dc260b58fe0e1a32b5c0b65 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Mon, 13 Jul 2026 15:37:51 -0700 Subject: [PATCH 07/14] drop no-op bridge fp8_model_init wrap and trim fp8 frozen base comments --- miles/backends/megatron_utils/actor.py | 6 ++---- miles/backends/megatron_utils/fp8_frozen_base.py | 14 +++++--------- miles/backends/megatron_utils/model.py | 2 +- 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 681a8ecead9..d4ba1a75ad3 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -152,9 +152,7 @@ def backup_source(): return source from miles.backends.megatron_utils.fp8_frozen_base import frozen_fp8_param_ids - # The frozen base lives as fp8 buffers; its bf16 param.data is transient - # (materialized by the forward pre-hook, freed at offload), so it must not - # be backed up or restored. + # the fp8-stored base's bf16 param.data is transient — never back up/restore it skip = frozen_fp8_param_ids(self.model) return ((name, tensor) for name, tensor in source if id(tensor) not in skip) @@ -218,7 +216,7 @@ def sleep(self) -> None: if not self.args.offload_train: return - if getattr(self.args, "fp8_frozen_base_store", False): + if self.args.fp8_frozen_base_store: from miles.backends.megatron_utils.fp8_frozen_base import free_frozen_base free_frozen_base(self.model) diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index 6346dd1e07c..27a9baf3126 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -71,14 +71,12 @@ def family(name: str) -> str: def install_fp8_hooks(module: torch.nn.Module) -> None: - # Materialize the bf16 weight before every forward. It is NOT freed after the - # forward: TE grouped-expert backward reads self.weight directly, and under - # activation recompute the recomputed forward + its backward straddle a free. - # The transient bf16 is released together at offload time (free_frozen_base). + # No post-forward free: TE grouped backward and activation recompute read self.weight + # after the forward, so the bf16 stays until free_frozen_base at offload. def pre(mod, inputs): for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): - q = getattr(mod, f"fp8q_{leaf}").contiguous() - s = getattr(mod, f"fp8s_{leaf}").contiguous() + q = getattr(mod, f"fp8q_{leaf}") + s = getattr(mod, f"fp8s_{leaf}") param.data = weight_dequant(q, s, BLOCK).to(dtype).reshape(shape) module.register_forward_pre_hook(pre) @@ -94,9 +92,7 @@ def frozen_fp8_param_ids(model_chunks) -> set[int]: def free_frozen_base(model_chunks) -> None: - """Drop the transient bf16 base weights so the offloaded/resident footprint is - the fp8 buffers. Call after the training step's backward (e.g. before offload). - The forward pre-hook re-materializes them on the next step.""" + """Drop the transient bf16 base so only the fp8 buffers stay resident; call post-backward.""" for chunk in model_chunks: for module in chunk.modules(): entries = getattr(module, "fp8_frozen_entries", None) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 140e8c9605f..6801df945bc 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -816,7 +816,7 @@ def initialize_model_and_optimizer( checkpointing_context={}, skip_load_to_model_and_opt=False, ) - if getattr(args, "fp8_frozen_base_store", False): + if args.fp8_frozen_base_store: from miles.backends.megatron_utils.fp8_frozen_base import quantize_frozen_base_to_fp8 quantize_frozen_base_to_fp8(model, args) From 31d05485f594b4c88c5bda3875b7f71b6407cabc Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Mon, 13 Jul 2026 15:44:39 -0700 Subject: [PATCH 08/14] enable fp8 frozen base store in the fp8 MoE-LoRA example --- examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh index 0a22c9d35ce..c7b5ed38641 100755 --- a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh +++ b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh @@ -110,6 +110,7 @@ PERF_ARGS=( --bf16 --fp8-format e4m3 --fp8-recipe blockwise + --fp8-frozen-base-store ) GRPO_ARGS=( From b63fbc7ded111cedefc39685b3b7f3b586952881 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Tue, 14 Jul 2026 13:32:19 -0700 Subject: [PATCH 09/14] bridge-LoRA: forward recompute args to the provider (port upstream #1593) --- miles/backends/megatron_utils/bridge_lora_helpers.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py index 1aa4293e127..00c09a788c8 100644 --- a/miles/backends/megatron_utils/bridge_lora_helpers.py +++ b/miles/backends/megatron_utils/bridge_lora_helpers.py @@ -93,6 +93,11 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: provider.sequence_parallel = args.sequence_parallel provider.virtual_pipeline_model_parallel_size = args.virtual_pipeline_model_parallel_size provider.context_parallel_size = args.context_parallel_size + provider.recompute_granularity = args.recompute_granularity + provider.recompute_method = args.recompute_method + provider.recompute_num_layers = args.recompute_num_layers + provider.recompute_modules = args.recompute_modules + provider.distribute_saved_activations = args.distribute_saved_activations provider.variable_seq_lengths = True provider.moe_token_dispatcher_type = "alltoall" provider.moe_router_load_balancing_type = "none" From 2ae12f04837308a1273295552b82c63c80bfe2f8 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Tue, 14 Jul 2026 15:06:20 -0700 Subject: [PATCH 10/14] deslop: drop dead plan-doc reference, record open TODOs in fp8_frozen_base --- examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh | 1 - miles/backends/megatron_utils/fp8_frozen_base.py | 5 +++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh index c7b5ed38641..44ce3589ab1 100755 --- a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh +++ b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh @@ -9,7 +9,6 @@ # - Block-wise FP8 e4m3 forward, BF16 backward + master weights # - --use-tis for MoE numerical drift compensation # -# See docs/superpowers/plans/2026-06-08-fp8-moe-lora-02-fp8-moe-lora-bringup.md. # for rerun the task pkill -9 sglang diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index 27a9baf3126..ba8607a50e9 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -10,6 +10,11 @@ logger = logging.getLogger(__name__) +# TODO: free the transient bf16 per layer during the step (peak win, ~base size at +# 122B); blocked on stale-weight retention across re-materializations. +# TODO: TE-native fp8 params (Float8BlockScaling) so GEMMs consume the stored fp8 +# directly and the bf16 copy never exists. + BLOCK = 128 # Frozen base linear weights to store as block-fp8 (mirrors the export-side From 7371cd399574c2fda09b2bdfdc4d5f52d3bbf33b Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Tue, 14 Jul 2026 12:07:56 -0700 Subject: [PATCH 11/14] fp8 frozen base: opt-in per-layer free to drop the rematerialization peak --- .../megatron_utils/fp8_frozen_base.py | 31 ++++++++++++++++--- miles/utils/arguments.py | 10 ++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index ba8607a50e9..59405991e11 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -75,17 +75,40 @@ def family(name: str) -> str: return "other" -def install_fp8_hooks(module: torch.nn.Module) -> None: - # No post-forward free: TE grouped backward and activation recompute read self.weight - # after the forward, so the bf16 stays until free_frozen_base at offload. +def install_fp8_hooks(module: torch.nn.Module, per_layer_free: bool) -> None: def pre(mod, inputs): for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): + if param.data.numel(): + continue q = getattr(mod, f"fp8q_{leaf}") s = getattr(mod, f"fp8s_{leaf}") param.data = weight_dequant(q, s, BLOCK).to(dtype).reshape(shape) module.register_forward_pre_hook(pre) + if not per_layer_free: + # bf16 stays until free_frozen_base at offload: TE grouped backward and + # activation recompute read self.weight after the forward. + return + + def free_entries(mod): + for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): + param.data = torch.empty(0, dtype=dtype, device=param.data.device) + + def post(mod, inputs, output): + # Under no_grad (checkpointed outer forward, forward-only logprob) nothing + # holds the weight for backward — the recompute pre-hook re-materializes it. + # Under grad (recompute replay, non-checkpointed) the backward still reads + # it, so free only once dgrad has reached the module inputs. + if not torch.is_grad_enabled(): + free_entries(mod) + return + grad_inputs = [t for t in inputs if isinstance(t, torch.Tensor) and t.requires_grad] + if grad_inputs: + torch.autograd.graph.register_multi_grad_hook(grad_inputs, lambda grads: free_entries(mod)) + + module.register_forward_hook(post) + def frozen_fp8_param_ids(model_chunks) -> set[int]: return { @@ -139,7 +162,7 @@ def quantize_frozen_base_to_fp8(model_chunks, args) -> None: if entries: module.fp8_frozen_entries = entries - install_fp8_hooks(module) + install_fp8_hooks(module, args.fp8_frozen_base_per_layer_free) if rank0(): logger.info( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2dd2e71bc74..e58f9dea4ca 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -220,6 +220,16 @@ def add_train_arguments(parser): "(dequant per layer in the forward) to halve resident base weight memory." ), ) + parser.add_argument( + "--fp8-frozen-base-per-layer-free", + action="store_true", + default=False, + help=( + "With --fp8-frozen-base-store: free each module's dequantized bf16 as soon " + "as the step no longer needs it (post-forward under no_grad, post-dgrad " + "otherwise) so train peak drops by ~the base size instead of holding it." + ), + ) parser.add_argument( "--allgather-cp", action="store_true", From df52d82bbb109118a3790bcaf8bb10f5e0154ca0 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Tue, 14 Jul 2026 17:01:44 -0700 Subject: [PATCH 12/14] fp8 frozen base: free via grad_fn post-hook, multi_grad_hook leaked whole graphs --- miles/backends/megatron_utils/fp8_frozen_base.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index 59405991e11..092c47f3ce6 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -10,8 +10,6 @@ logger = logging.getLogger(__name__) -# TODO: free the transient bf16 per layer during the step (peak win, ~base size at -# 122B); blocked on stale-weight retention across re-materializations. # TODO: TE-native fp8 params (Float8BlockScaling) so GEMMs consume the stored fp8 # directly and the bf16 copy never exists. @@ -98,14 +96,18 @@ def free_entries(mod): def post(mod, inputs, output): # Under no_grad (checkpointed outer forward, forward-only logprob) nothing # holds the weight for backward — the recompute pre-hook re-materializes it. - # Under grad (recompute replay, non-checkpointed) the backward still reads - # it, so free only once dgrad has reached the module inputs. + # Under grad (recompute replay) free after the module's backward: TE forwards + # are a single autograd Function, so output.grad_fn is the node that reads the + # weight; its post-hook fires once dgrad is done (frozen base has no wgrad). + # NOTE: multi_grad_hook is unusable here — its closure holds the inputs' + # grad_fns from hooks inside the same graph, an uncollectible cycle that + # retained every microbatch graph (~14 GB/step). if not torch.is_grad_enabled(): free_entries(mod) return - grad_inputs = [t for t in inputs if isinstance(t, torch.Tensor) and t.requires_grad] - if grad_inputs: - torch.autograd.graph.register_multi_grad_hook(grad_inputs, lambda grads: free_entries(mod)) + out = output[0] if isinstance(output, tuple) else output + if isinstance(out, torch.Tensor) and out.grad_fn is not None: + out.grad_fn.register_hook(lambda grad_inputs, grad_outputs: free_entries(mod)) module.register_forward_hook(post) From b5578b9f398a3b05144835bca09bf6d5a6cda89f Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Tue, 14 Jul 2026 20:31:34 -0700 Subject: [PATCH 13/14] enable per-layer free in the fp8 MoE-LoRA example --- examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh index 44ce3589ab1..c1931201890 100755 --- a/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh +++ b/examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh @@ -110,6 +110,7 @@ PERF_ARGS=( --fp8-format e4m3 --fp8-recipe blockwise --fp8-frozen-base-store + --fp8-frozen-base-per-layer-free ) GRPO_ARGS=( From 63ddd29246a62daec79b6fc11c511a6f34e57356 Mon Sep 17 00:00:00 2001 From: MuuSeoTia Date: Wed, 15 Jul 2026 11:26:23 -0700 Subject: [PATCH 14/14] deslop fp8 frozen base: drop redundant checks, tighten hook comments --- .../megatron_utils/fp8_frozen_base.py | 28 ++++++++----------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/miles/backends/megatron_utils/fp8_frozen_base.py b/miles/backends/megatron_utils/fp8_frozen_base.py index 092c47f3ce6..ff02c92428a 100644 --- a/miles/backends/megatron_utils/fp8_frozen_base.py +++ b/miles/backends/megatron_utils/fp8_frozen_base.py @@ -56,7 +56,7 @@ def is_base_linear_weight(name: str) -> bool: def should_quantize(name: str, param: torch.nn.Parameter) -> bool: - return param is not None and param.dim() == 2 and not param.requires_grad and is_base_linear_weight(name) + return param.dim() == 2 and not param.requires_grad and is_base_linear_weight(name) def family(name: str) -> str: @@ -90,18 +90,15 @@ def pre(mod, inputs): return def free_entries(mod): - for leaf, (param, shape, dtype) in mod.fp8_frozen_entries.items(): + for param, _, dtype in mod.fp8_frozen_entries.values(): param.data = torch.empty(0, dtype=dtype, device=param.data.device) def post(mod, inputs, output): - # Under no_grad (checkpointed outer forward, forward-only logprob) nothing - # holds the weight for backward — the recompute pre-hook re-materializes it. - # Under grad (recompute replay) free after the module's backward: TE forwards - # are a single autograd Function, so output.grad_fn is the node that reads the - # weight; its post-hook fires once dgrad is done (frozen base has no wgrad). - # NOTE: multi_grad_hook is unusable here — its closure holds the inputs' - # grad_fns from hooks inside the same graph, an uncollectible cycle that - # retained every microbatch graph (~14 GB/step). + # no_grad (checkpointed outer forward, forward-only): nothing saves the weight, + # free now — the pre-hook re-materializes on recompute. Under grad, free after + # the module's backward via its grad_fn post-hook (TE forward = one autograd + # Function; frozen base has no wgrad). multi_grad_hook cannot be used here: + # its closure↔graph reference cycle retains every microbatch graph. if not torch.is_grad_enabled(): free_entries(mod) return @@ -125,10 +122,7 @@ def free_frozen_base(model_chunks) -> None: """Drop the transient bf16 base so only the fp8 buffers stay resident; call post-backward.""" for chunk in model_chunks: for module in chunk.modules(): - entries = getattr(module, "fp8_frozen_entries", None) - if not entries: - continue - for leaf, (param, shape, dtype) in entries.items(): + for param, _, dtype in getattr(module, "fp8_frozen_entries", {}).values(): param.data = torch.empty(0, dtype=dtype, device=param.data.device) @@ -166,11 +160,11 @@ def quantize_frozen_base_to_fp8(model_chunks, args) -> None: module.fp8_frozen_entries = entries install_fp8_hooks(module, args.fp8_frozen_base_per_layer_free) - if rank0(): + if rank0() and counts: logger.info( - "fp8 frozen base: quantized %d tensors %s, freed ~%.2f GB/rank, roundtrip_relerr=%s", + "fp8 frozen base: quantized %d tensors %s, freed ~%.2f GB/rank, roundtrip_relerr=%.2e", sum(counts.values()), counts, freed_bytes / (1024**3), - f"{roundtrip_relerr:.2e}" if roundtrip_relerr is not None else "n/a", + roundtrip_relerr, )