Repository navigation
Tracking: numerical issues found in an audit of ART main after #1105 #1111
Description
After #1105 fixed a BF16 cross-entropy dead zone in TrainerRank, we audited the rest of ART at e2a406a for numerical issues. This issue tracks what we found. The audit found 11 items. The top three have dedicated issues; the other eight are described below in enough detail to open dedicated issues later. Items are ranked by expected impact × likelihood on live training paths. Line links point to e2a406a. Every item listed here also exists in the revision before #925.
Dedicated issues
- Megatron backend: Transformer Engine fused cross-entropy zeroes target gradients for p ≥ 0.998 (same bug as #1105) #1108: Megatron backend: Transformer Engine's fused cross-entropy zeroes target gradients for p ≥ 0.998 (the same bug as fix(trainer-rank): combine cross-entropy target and softmax gradients in FP32 #1105). Affects RL and SFT. Confirmed by running TE's own kernel.
- Unsloth backend: BF16 logsumexp zeroes logprobs and target gradients for p ≥ ~0.88 #1109: Unsloth backend: BF16
logsumexpzeroes logprobs and target gradients for p ≥ about 0.88. This is on ART's default local backend. Confirmed on CPU with the same operations; the GPU dtype still needs confirming. - GRPO: equal-reward groups can become a full-scale advantage signal or NaN after pack normalization #1110: GRPO: equal-reward groups can become a full-scale advantage signal, or NaN, after pack normalization. Confirmed.
Not yet opened
4. BF16 LM-head output logits (TrainerRank, Megatron, serving)
- Where:
trainer_rank/_impl.py#L5948-L5971(_local_logits_from_hidden_rows). The Megatron output layer and vLLM do the same. - What: logits come out of the output-layer GEMM in BF16. After fix(trainer-rank): combine cross-entropy target and softmax gradients in FP32 #1105, everything downstream is FP32.
- Why it might matter: logits around 35–47 have a BF16 spacing of 0.25, so each logit carries up to ±0.125 of rounding.
- Measured on CPU with Qwen3-4B (294 assistant tokens):
- per-token logprob error is 0.075 nats at p90 and 0.13 at p99;
- the summed hidden-state gradient changes by 1.1% (cosine 0.99994);
- the expected bias of 1 − p_t is about +0.5%.
- Effect: noise rather than bias for SFT. For RL, roughly 0.1 nats of per-token noise matters next to a PPO clip of ε = 0.2, and it may add to the trainer/sampler mismatch.
- Confidence: confirmed by measurement. Impact is likely low for SFT and plausible-to-moderate for RL ratios.
- Suggested check: a GPU A/B with an FP32-output GEMM for the head chunk. Compare trainer and vLLM logprobs on the same tokens: the |Δ| quantiles and the fraction with |Δ| > 0.1. Admission must price the larger chunk.
5. TrainerRank accumulates LoRA gradients in BF16
- Where:
trainer_rank/_versions.py#L43-L51and#L270-L274. The upcast happens only atoptim_step. - What: gradients from each backward call are summed in the parameter dtype (BF16). The Megatron backend instead accumulates in an FP32
main_grad(grad_reduce_in_fp32=True). - Measured against an FP32 accumulator: 0.41% relative L2 error at 16 micro-batches per step and 1.8% at 256 (1 − cosine ≈ 1e-5). The error is unbiased.
- Effect: small noise next to SGD noise, but inconsistent with the Megatron backend.
- Confidence: confirmed by repro.
- Suggested fix: accumulate into FP32 staging tensors, or keep FP32 LoRA leaves (LoRA memory is small).
6. TrainerRank's Adam epsilon isn't configurable and differs from the Megatron backend
- Where:
trainer_rank/_optimizer.py#L423-L428.AdamParamsat_impl.py#L120-L126has noeps, so PyTorch's 1e-8 default applies. The Megatron backend usesadam_eps=1e-13(megatron/train.py#L283). - Why it might matter: for tiny gradient elements, such as LoRA A shortly after B leaves 0 with tight clipping, 1e-8 is close to √v and damps updates.
- Effect: a small difference in effective learning rate early in training between backends.
- Confidence: confirmed by reading; impact is low.
- Suggested fix: add
epstoAdamParamsand the cross-rank config guard.
7. Truncated importance sampling falls back to the policy ratio
- Where:
loss.py#L375-L385. - What: with
truncated_importance_samplingset and nooriginal_logprobs, the loss is multiplied byclamp(prob_ratio, max=U), whereprob_ratiois still π_new/π_old (possibly clamped, masked or sequence-level). - Why it might matter: that's not the trainer/sampler correction TIS is meant to apply. The policy ratio gets applied a second time, and CISPO already weights by the clipped ratio.
- Effect: biased RL weighting whenever
original_logprobsare missing. - Affects: RL with TIS enabled on the Megatron, Unsloth and legacy Tinker backends.
- Confidence: plausible (logic).
- Suggested fix: skip TIS, or raise, when
original_logprobsis absent.
8. Unclamped exp(new - old) on the PPO path
- Where:
loss.py#L299and#L363-L367. - What: if an old logprob is −inf, a sentinel such as −9999, or extremely stale, the ratio is inf. With
ppo=Trueand A < 0,min(inf·A, clip·A)is −inf, so the loss is +inf and the gradients are NaN. CISPO's detached clip is safe. - Effect: a NaN step. On Megatron,
_validate_train_step_result_finitecatches it later. - Confidence: speculative, since vLLM's raw logprobs at T = 1 are finite.
- Suggested fix: clamp the log-ratio (for example to ±20) before
exp, or assert finite old logprobs on trainable tokens.
9. The legacy Tinker service builds the wrong per-datum masks and steps per packed row
tinker/service.py#L136-L180.- Masks are built as
(group == g) | (parent == parent(g)). Prefix-tree roots haveparent_id == group_id(megatron/prefix_tree_packing.py#L258), so sibling branches leak into each datum's context, and the root group's datum covers the whole row. optim_step_asyncsits inside the per-row loop.
(group == g) | (group == parent(g)), or train unpacked. Accumulate forward_backward over rows and step once.10. tinker_native uses the sample standard deviation for advantages
- Where:
tinker_native/data.py#L32-L44, which usestorch.std(n−1). The other backends use the population standard deviation (n) inpreprocessing/tokenize.py#L763-L771. - Effect: advantages differ by a constant factor of √((n−1)/n), which is 0.87 at n = 4. That changes the effective learning rate, clip and KL scale between backends.
- Confidence: confirmed by reading; impact is low.
- Suggested fix: pick one convention.
11. The LoRA delta is added in BF16 in the forward pass
- Where:
megatron/lora.py#L1019-L1020and#L1414. - What:
(x @ A) @ B * sis computed in BF16 and added to the BF16 base output. While the delta is below half a BF16 ulp of the base activation (early in training), part of it rounds away. Gradients are unaffected. - Effect: the forward loss moves in quantized steps early on. This matches vLLM, which also adds in BF16.
- Confidence: speculative; inherent to BF16 LoRA.
- Suggested action: probably leave it. To test, compare the one-step logprob change with the delta added in FP32.
Checked and found sound
- TrainerRank's head after fix(trainer-rank): combine cross-entropy target and softmax gradients in FP32 #1105:
- Triton stats and backward kernels keep max and sum in FP32 and combine softmax, top-k and target gradients in FP32 before a single store.
- The eager and unchunked fallbacks are FP32.
- The TP log-normalizer is FP32, and ignored labels and unowned TP columns are masked correctly.
- The CP row gather routes one gradient copy to each owner.
- TrainerRank's optimizer and checkpoints:
- FP32 master weights and FP32 Adam moments are used.
scale_gradsis applied in FP32 before the DP/CP sum, and TP gradient sync runs in FP32.- The FP32 grad norm counts replicas once, and clipping uses
min(1, c / (‖g‖ + 1e-6)). - Checkpoints store the FP32 state.
- Importance-sampling corrections (
_corrections.py): computed in FP32/FP64 log space, clamped beforeexp, with non-finite inputs rejected. - Megatron-backend normalization:
- per-token loss uses the global token count over DP×CP, so there's no mean-of-means;
- micro-batch gradients accumulate in an FP32
main_grad; - the logprob and score helpers are FP32.
- Label alignment under prefix-tree packing and context parallelism: correct.
loss.pycore: the ratio, clip, CISPO and PPO math are FP32 on the Megatron and TrainerRank consumer paths. NaN old logprobs are replaced with ratio 1.- vLLM logprob transport: raw logprobs at temperature 1 travel as FP32 JSON floats, with no quantization.
- LoRA export: there's no B·A product on ART's side, the alpha round-trip is asserted, and tensors are exported in the parameter dtype.
- Spot-checked kernels: CP attention merge (FP32 LSE with −inf-safe
logaddexp), GDN gate and state math (FP32), and the grouped expert LoRA (FP32 accumulation).
Activity
fix: bound policy importance ratios before exponentiation
fix: isolate legacy Tinker branch contexts
fix: use population std for native Tinker advantages
fix: skip TIS when original logprobs are unavailable
fix: bound policy importance ratios before exponentiation
fix: isolate legacy Tinker branch contexts
fix: use population std for native Tinker advantages
fix: skip TIS when original logprobs are unavailable
fix: bound policy importance ratios before exponentiation
fix: isolate legacy Tinker branch contexts
fix: use population std for native Tinker advantages
fix: skip TIS when original logprobs are unavailable