Fix two FP8-training issues: PLE ngram dequantization and LoRA fp8 autocast guard - #183
Open
ehooon wants to merge 2 commits into
Open
Fix two FP8-training issues: PLE ngram dequantization and LoRA fp8 autocast guard #183ehooon wants to merge 2 commits into
ehooon wants to merge 2 commits into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two independent fixes that only manifest once FP8 enters training; bf16 workflows are
unaffected.