`ValueError: can't optimize a non-leaf Tensor` when using `--optimizer-cpu-offload` (HybridDeviceOptimizer) together with `mark_keep_in_fp32` parameters (e.g. DeepSeek V4 CSA `ape`/`attn_sink`, mHC `alpha_*`/`bias`)

[HowardZorn]

Describe the bug

When --optimizer-cpu-offload is enabled (i.e. the base optimizer is wrapped by HybridDeviceOptimizer) together with --use-distributed-optimizer, training crashes at optimizer construction time with:

ValueError: can't optimize a non-leaf Tensor

Root cause: DistributedOptimizer._build_model_and_main_param_groups() in megatron/core/optimizer/distrib_optimizer.py builds the "shard" tensor that is registered with the inner optimizer differently depending on the model parameter's dtype:

# bf16 / fp16 branch (line ~389) — calls .detach() first, safe
shard_model_param = model_param.detach().view(-1)[param_range.start : param_range.end]

# fp32 branch (line ~474) — does NOT call .detach(), buggy
elif model_param.type() == 'torch.cuda.FloatTensor':
    shard_model_param = model_param.view(-1)[param_range.start : param_range.end]

Because model_param.requires_grad is always True (asserted a few lines above), shard_model_param in the fp32 branch is a non-leaf tensor with requires_grad=True (a ViewBackward op on a tensor that requires grad).

When the base optimizer is a HybridDeviceOptimizer (--optimizer-cpu-offload), DistributedOptimizer.__init__ instead re-constructs it by calling its __init__ again:

if isinstance(self.optimizer, HybridDeviceOptimizer):
    self.optimizer = HybridDeviceOptimizer(
        params=[g["orig_group"] for g in self.opt_group_ranges], **self.optimizer.defaults
    )

HybridDeviceOptimizer.__init__ calls torch.optim.Optimizer.__init__ → add_param_group(), which does validate that every parameter is a leaf tensor (see torch/optim/optimizer.py), and raises the ValueError above.

This only manifests when the model actually contains at least one real torch.cuda.FloatTensor trainable parameter that goes through the shard_fp32_params_this_group path. In this codebase that happens whenever a parameter is intentionally kept in FP32 via megatron.core.transformer.module.mark_keep_in_fp32, e.g.:

• ape / attn_sink in DeepSeek V4 sparse attention (megatron/core/transformer/experimental_attention_variant/csa.py), explicitly called out in the mark_keep_in_fp32 docstring.
• mapping_proj.weight, alpha_pre, alpha_post, alpha_res, bias in the mHC (Manifold-Constrained Hyper-Connections) module (megatron/core/transformer/hyper_connection.py, enable_hyper_connections=True).
• hc_head_fn / hc_head_base / hc_head_scale in megatron/core/transformer/multi_token_prediction.py and megatron/core/models/hybrid/hybrid_block.py.

These parameters are deliberately excluded from the FP16/BF16 cast by convert_module_to_dtype_except_fp32_marked() inside Float16Module, so they remain torch.cuda.FloatTensor end-to-end — which is exactly the type that hits the un-.detach()-ed branch above.

Dense models without any mark_keep_in_fp32 parameters (e.g. plain Llama/Qwen-style GPT) never create a real fp32 model parameter, so they neverexercise this code path and never hit the bug — regardless of --optimizer-cpu-offload. This is why the crash is only observed with models such as DeepSeek V4 (CSA sparse attention / mHC hyper-connections).

Steps/Code to reproduce bug

1. Minimal, non-distributed repro of the underlying PyTorch-level defect (isolates exactly what HybridDeviceOptimizer.__init__ hits):

   import torch
   
   class Dummy(torch.optim.Optimizer):
       def __init__(self, params):
           super().__init__(params, defaults={})
   
   # Simulates a FP32 model param (e.g. CSA `ape`, mHC `alpha_pre`) that is
   # requires_grad=True, as built by _build_model_and_main_param_groups()
   # WITHOUT .detach() in the fp32 branch.
   p = torch.nn.Parameter(torch.zeros(4))
   shard = p.view(-1)[0:2]
   print(shard.requires_grad, shard.is_leaf)  # True False
   
   Dummy([{'params': [shard]}])
   # ValueError: can't optimize a non-leaf Tensor

2. Full end-to-end repro: pretrain a GPT model that has at least one mark_keep_in_fp32 parameter (e.g. DeepSeek V4 with CSA sparse attention, or any model with --enable-hyper-connections) with:

   --use-distributed-optimizer
   --optimizer-cpu-offload
   --optimizer-offload-fraction 1.0
   --use-precision-aware-optimizer
   

   Training fails immediately during MegatronTrainer construction (get_megatron_optimizer → _get_megatron_optimizer_based_on_param_groups → DistributedOptimizer.__init__ → HybridDeviceOptimizer.__init__ → torch.optim.Optimizer.add_param_group) with:

   File ".../megatron/core/optimizer/distrib_optimizer.py", line 754, in __init__
       self.optimizer = HybridDeviceOptimizer(
   File ".../megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py", line 57, in __init__
       super(HybridDeviceOptimizer, self).__init__(
   File ".../torch/optim/optimizer.py", line 408, in __init__
       self.add_param_group(cast(dict, param_group))
   File ".../torch/_dynamo/eval_frame.py", line 1263, in _fn
       return fn(*args, **kwargs)
   File ".../torch/optim/optimizer.py", line 1159, in add_param_group
       raise ValueError("can't optimize a non-leaf Tensor")
   ValueError: can't optimize a non-leaf Tensor
   

Expected behavior

--optimizer-cpu-offload should work regardless of whether the model has mark_keep_in_fp32 (FP32-pinned) parameters. _build_model_and_main_param_groups() should build the fp32 shard the same, safe way as the bf16/fp16 branch, i.e. call .detach() before .view():

elif model_param.type() == 'torch.cuda.FloatTensor':
    shard_model_param = model_param.detach().view(-1)[param_range.start : param_range.end]
    ...

We verified locally that this one-line fix (mirroring the existing bf16/fp16 branch) resolves the crash and training proceeds normally with --optimizer-cpu-offload enabled. .detach() only removes the tensor from the autograd graph (requires_grad: True → False, is_leaf: False → True); it does not change dtype or storage, and Megatron already manages .grad / .data.copy_() on these shard tensors manually elsewhere, so this should be a safe, non-invasive fix.

Additional context

• megatron-core version: 0.19.0.dev (repo on dev branch, commit d8b71082ea7033974703a101ded49a359b773458)
• torch: 2.11.0+cu130
• transformer_engine: 2.14.1+366798ef
• Affected file: megatron/core/optimizer/distrib_optimizer.py, DistributedOptimizer._build_model_and_main_param_groups() (fp32 branch around line 474 on dev)
• Relevant training args from the failing run: optimizer='adam', optimizer_cpu_offload=True, optimizer_offload_fraction=1.0, use_precision_aware_optimizer=True, use_distributed_optimizer=True, pipeline_model_parallel_size=8, virtual_pipeline_model_parallel_size=3, expert_model_parallel_size=4, model = DeepSeek-V4-Flash (multi_latent_attention=True, experimental_attention_variant='dsv4_hybrid', qk_layernorm=True, enable_hyper_connections=True).
• Only the ranks holding the pipeline stage(s) whose parameters include a mark_keep_in_fp32-marked tensor crash; other ranks are unaffected.

[ilml]

@deepakn94 could you plz take a look?