[StevenWang-CY]

Motivation

specforge export --to sglang serializes the state of a draft that materialize_draft builds in BF16, so a checkpoint trained with model.torch_dtype: float32 or float16 is rounded on export. --to hf has preserved the checkpoint tensors since #872, but both exporters still stamp dtype: bfloat16 into the exported config.json, and Transformers 5.12.1 resolves the reload dtype from that config: AutoDraftModel.from_pretrained(export_dir) without an explicit dtype casts the preserved FP32 weights back to BF16. The export test module's docstring promises bit-identical round trips for both formats, and the existing CUDA gate only checks that with a BF16-trained draft, where the cast is a no-op.

Modifications

• materialize_draft builds the draft in the checkpoint's floating-point dtype (checkpoint_dtype, the single dtype shared by the draft weights), so the module's parameters and the dtype save_pretrained stamps into the config both carry the trained precision. A checkpoint that mixes dtypes keeps the BF16 fallback.
• export_to_sglang serializes the checkpoint tensors as stored, exactly as export_to_hf does, while keeping the t2d/d2t buffers loaded from an explicit --vocab-mapping override. _serving_state validation and weight-map renames are unchanged.
• BF16 checkpoints, the schema default, export byte-for-byte as before: the same config.json and model.safetensors hashes on unmodified main and on this head for both formats.
• The external-embedding CLI test from fix(export): honor explicit vocabulary mappings in HF exports #872 now expects the embedding in the export dtype, which follows the checkpoint (FP32 there) instead of a fixed BF16.
• One documentation paragraph in docs/sections/basic_usage/training.md states the precision contract.

Serving is unaffected: SGLang casts draft weights to the server dtype at load time. FP32 checkpoints produce larger export files than before, at the precision they were trained in.

Related Issues

Follow-up to #872. No linked issue; the live audit of open PRs and issues found no competing precision fix. #811 (draft) and #882 touch neighbouring lines in these files for MoE and expert-parallel layouts, not precision.

Accuracy Test

Validation on head 6f8629588a1aaf8c27811ab9b14957a9b9cdf190 against unmodified main ed64d275bac8a48e126adb2827368603566a3029, CPU only (macOS, torch 2.13.0, transformers 5.12.1; no CUDA, sglang, or triton available here):

• New TestExportCheckpointPrecision (three CPU tests, FP32/FP16/BF16 subtests) fails on unmodified main for the intended reasons: torch.bfloat16 != torch.float32/float16 on every serialized SGLang weight, and config dtype 'bfloat16' != 'float32' on the HF export; the BF16 subtest passes on main as the control. All pass on this head.
• python -m unittest tests.test_runtime.test_export: 18 tests, 4 CUDA-only skips, all passing.
• python -m unittest discover -s ./tests -p "test_*.py" in a fresh clean worktree of this exact commit: 1158 tests; the 17 failures and 20 errors are identical to the unmodified-main baseline in this environment and are all environmental (no CUDA, no sglang, no triton).
• pre-commit run --all-files passes.
• BF16 byte-identity control: exporting the same BF16 checkpoint through export_to_sglang and export_to_hf on main and on this head yields identical SHA-256 for config.json and model.safetensors.

The hosted PR Test job on this head has since run the CUDA-gated export round-trips that CPU runs skip — test_to_hf_reloads_bit_identical, test_to_sglang_produces_exact_serving_keys, test_weight_map_renames_are_applied — alongside the three new precision tests, and all of them passed; lint passed too.

That job's only error is unrelated to this change: test_dp_reduction_and_ragged_shards failed in dist.init_process_group with EADDRINUSE on port 46497, the rendezvous port its own _free_port() helper binds, closes, and then hands to mp.spawn. It ran before any test in this PR's module, it passes locally with this patch, and the same runner passed it on later runs of other branches. A rerun clears it.

Serving is unaffected by the file's precision: SGLang v0.5.18's default_weight_loader copies each tensor into the server-dtype parameter (param.data.copy_(loaded_weight)), and its EAGLE3 draft loader reads d2t/t2d as integer buffers. As the export docs already state, loading a draft in a real speculative-decoding server and measuring acceptance remains a GPU serving validation step.

Benchmark & Profiling

Not applicable; export-time code path only.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.

[StevenWang-CY]

The unit-test job on this head (6f86295) ran 1249 tests with one error and no failures. The error is test_runtime.test_evaluator_aggregation.TestEvaluatorDataParallel.test_dp_reduction_and_ragged_shards, which spawns a two-rank gloo group on a port obtained from _free_port() (bind port 0, close, reuse later) and hit EADDRINUSE on port 46497 when the workers bound it: job log.

That test imports nothing from specforge/export, this PR changes only the two exporter modules, their test module and one docs paragraph, and the same test passes on the other recent PR Test runs, so the error is the rendezvous port being retaken between the free-port probe and the workers' bind on the shared runner rather than a change in behaviour. The export tests added here passed in the same job. A rerun should be green; I can also push a rebase if maintainers prefer a fresh run.

[StevenWang-CY]

@maocheng23, could you review this exporter follow-up to #872? It keeps the SGLang export at the checkpoint's precision and makes the exported config.json dtype match the serialized tensors, so an HF reload without an explicit dtype no longer casts preserved FP32 weights back to BF16.

On head 6f86295: lint passes, and the unit-test job ran 1249 tests with the new export tests green and a single unrelated error, test_dp_reduction_and_ragged_shards, from an EADDRINUSE on its free-port rendezvous (details in the comment above). A rerun of that job should be green. BF16 checkpoints export byte-for-byte as before, and I checked against SGLang v0.5.18's default_weight_loader, which copies into the server-dtype parameter, so FP32 safetensors serve unchanged.

[StevenWang-CY]

@FrankLeeeee, ready for review; CODEOWNERS assigned you via /tests.

Head 6f86295: lint green. The PR Test job ran the CUDA-gated export round-trips (test_to_hf_reloads_bit_identical, test_to_sglang_produces_exact_serving_keys, test_weight_map_renames_are_applied) plus the three new precision tests, all passing.

The job's one error, test_dp_reduction_and_ragged_shards, is unrelated. It fails in dist.init_process_group with EADDRINUSE on the rendezvous port its own _free_port() helper binds, closes, and hands to mp.spawn. It runs before any test this PR touches, passes locally with the patch, and passes on the same runner for other branches. A rerun clears it.

[jiapingW]

@StevenWang-CY Hi, I have reviewed it and add some new tests. I think the loading logic will need to be updated to accommodate mixed-precision draft models later. For instance, the "dspark" variant of DeepSeek-V4-Flash-0731.