Skip to content

Fix two FP8-training issues: PLE ngram dequantization and LoRA fp8 autocast guard - #183

Open
ehooon wants to merge 2 commits into
modelscope:mainfrom
ehooon:fix-fp8-training
Open

Fix two FP8-training issues: PLE ngram dequantization and LoRA fp8 autocast guard #183
ehooon wants to merge 2 commits into
modelscope:mainfrom
ehooon:fix-fp8-training

Conversation

@ehooon

@ehooon ehooon commented Sep 3, 2026

Copy link
Copy Markdown

Two independent fixes that only manifest once FP8 enters training; bf16 workflows are
unaffected.

 1. **PLE ngram embedding loaded without dequantization.** FP8 checkpoints store                                            
    `ple.ple_embedding.ngram_embedding` as F8_E4M3 shards plus a scalar `weight_scale`                                      
    (unlike experts' blockwise `weight_scale_inv`). The to_mcore path copied raw fp8                                        
    values into the bf16 embedding without the scale, so the table was off by the full                                      
    scale factor (~5000x) and first-step loss was 12.6 ≈ ln(vocab) — effectively                                           
    random — instead of ~0.3. Fix: multiply by the scale in fp32 on load; symmetrically                                     
    divide and re-quantize to fp8 on to_hf export, writing the scale key back. bf16                                         
    checkpoints (no scale key) load as-is with a warning.                                                                   
                                                                                                                            
 2. **LoRA A/B crash under fp8 autocast.** With `fp8_param=True`, the global fp8                                            
    autocast also covers the rank-sized LoRA A/B weights (e.g. `[4, 2560]`), violating                                      
    TE's FP8 GEMM divisibility rules. Fix: wrap the LoRA branch of the TELinear path in                                     
    `fp8_autocast(enabled=False)` when `fp8_param` is set — LoRA runs in bf16, base                                         
    weights stay fp8 (mirrors the `in_proj_ba` guard in `gpts/qwen4_exp.py`).                                               
                                                                                                                            
 Adds a numerical round-trip unit test (`tests/test_qwen4_exp_ple_fp8_scale.py`,                                            
 synthetic tensors, no model download).                                                                                     
                                                                                                                            
 Verified by Megatron LoRA SFT smoke runs of Qwen3.8-Flash-Next-FP8 on 2 GPUs                                               
 (smoke config includes TP=2+SP, EP=2, PP=1, CP=1): bf16-dequantized first-step loss                                        
 0.28 (was 12.6); fp8-param mode 12/12 steps, converging normally. Smoke runs used                                          
 PP=1; since to_hf reuses the scale stashed on the bridge instance during to_mcore,                                         
 a same-process load→export round trip with PP>1 (PLE layer and saving rank on                                              
 different stages) is not covered.   

…FP8 checkpoints

FP8 checkpoints store the PLE ngram embedding table as F8_E4M3 shards plus
a single scalar weight_scale (true value = weight * scale), unlike experts
which use blockwise weight_scale_inv. The to_mcore path copied raw fp8
values into the bf16 embedding without applying the scale, making the table
off by an order of magnitude (first-step loss 12.6 instead of ~0.3).

This only affects training that starts from an FP8-format checkpoint; bf16
checkpoints have no weight_scale key and are loaded as-is (with a warning).

Multiply by the scale in fp32 on load, and symmetrically divide by the
stashed scale and cast back to fp8 on the to_hf export path, writing the
scale key back into the state dict.

Add a numerical round-trip unit test (synthetic tensors, no model needed).
This only triggers when training with fp8_param=True: the global
fp8_autocast then also covers the LoRA A/B projections, whose rank-sized
weights (e.g. [4, 2560]) violate TE's FP8 GEMM divisibility rules and crash
the forward pass.

Wrap the LoRA branch of the TELinear path in fp8_autocast(enabled=False)
when fp8_param is set, so the LoRA side path runs in bf16 while the base
weights stay in fp8 -- mirroring the existing in_proj_ba guard in the GDN
path of gpts/qwen4_exp.py.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant