feat(model): thread weight_dtype through HF export for plain-dtype DeepSeek-V4 output

[Meirtz]

What

Thread weight_dtype: Optional[torch.dtype] = None through the HF export path — export_hf_weights / save_hf_pretrained / stream_weights_megatron_to_hf — carried per-task via a new optional WeightConversionTask.weight_dtype field. When set, the DeepSeek-V4 bridge emits plain weights in that dtype (no *.scale companions) instead of re-creating the source repo's quantized layout. Default (None) keeps today's behavior. CLI: --export-weight-dtype on the export subcommand.

Why

DSv4 HF export unconditionally re-creates the source repo's quantized weight/scale layout (maybe_modify_converted_hf_weight → requantize_hf_weight_scale_pairs, from #3969). That's right for checkpoint conversion, but bf16-SFT'd weights get silently post-hoc quantized — a user found *.scale tensors in their SFT export and asked about train/inference parity.

Design (revised after reviewer feedback): the requantize hook runs on both export consumers — online weight streaming to rollout engines (export_hf_weights, e.g. verl RL weight sync) and on-disk checkpoints (save_hf_pretrained) — so a bridge-level boolean cannot configure them independently. A dtype-typed parameter on each public API lets callers choose per path (e.g. bf16 to rollout for RL parity, quantized to disk for serving-format artifacts, or vice versa). Hook signatures are unchanged (the dtype rides on the task), so the other bridges overriding this hook (dsv3, gemma4, kimi, mimo, flux) are unaffected; DSv3 can adopt the same field later.

Verified

• the saver streams only yielded tensors (omitting .scale keys is safe); exported config.json is built fresh (torch_dtype: bfloat16, no quantization fields); safetensors index regenerated from written tensors;
• unit tests cover dtype-set (plain weights, non-float tensors untouched) and default (requantize) paths.

Notes

• AI-assisted (Claude); analysis, validation and review by the human author.

🤖 Generated with Claude Code

[Meirtz]

Reworked per reviewer feedback (offline discussion): the hook serves two export consumers — online weight streaming to rollout engines and on-disk checkpoints — so the bridge-level boolean is gone. Now weight_dtype: Optional[torch.dtype] on export_hf_weights / save_hf_pretrained, carried per-task (WeightConversionTask.weight_dtype), hook signatures unchanged. d10a8e7e.

[Meirtz]

Full-model E2E validation (DeepSeek-V4-Flash, 43 layers, real weights, TP1/PP4/EP8 on 8×GB300; same imported Megatron checkpoint for both runs):

export	tensors	.scale tensors	dtypes
default (quantized)	69,187	34,167	F8_E4M3, F8_E8M0, I8 (MXFP4), BF16, F32, I32
--export-weight-dtype bfloat16	35,020	0	BF16, I32

35,020 + 34,167 = 69,187 — the bf16 artifact contains exactly every weight with no scale companions (I32 = tid2eid routing buffers, correctly left untouched).

Two notes from the run: (1) the smoke caught a real bug in the first version of this PR — WeightConversionTask is a frozen dataclass, so the dtype is now applied via dataclasses.replace (ce66e82), and the unit tests now use real task instances instead of mocks; (2) with weight_dtype set, the saver's completeness check counts the source's .scale entries as "not written" (export proceeds with --not-strict) — cosmetic, can be polished in a follow-up by excluding scale keys from the expected set.

[Meirtz]

/ok to test 1b93e3c

[cuichenx]

LGTM

[cuichenx]

please check comments above

[Meirtz]

/ok to test b577be3

[Meirtz]

/ok to test 5a97742

[HollowMan6]

Thank you @Meirtz ! Now the code logic looks much better and I don't have more comments