Tracking: numerical issues found in an audit of ART main after #1105

This is a tracking issue, so the work is to open the dedicated issues it lists, not to patch code. Start from the linked files: trainer_rank/_impl.py for the BF16 logits, trainer_rank/_versions.py for LoRA gradient accumulation, loss.py for the importance-sampling and PPO ratio paths, and tinker/service.py for the mask construction. Done means each unopened item has its own issue with the file, the expected behaviour and a reproduction or test.

After #1105 fixed a BF16 cross-entropy dead zone in TrainerRank, we audited the rest of ART at e2a406a48 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 e2a406a48. Every item listed here also exists in the revision before #925.

Dedicated issues

• #1108: Megatron backend: Transformer Engine's fused cross-entropy zeroes target gradients for p ≥ 0.998 (the same bug as #1105). Affects RL and SFT. Confirmed by running TE's own kernel.
• #1109: Unsloth backend: BF16 logsumexp zeroes 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.
• #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 #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-L51 and #L270-L274. The upcast happens only at optim_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. AdamParams at _impl.py#L120-L126 has no eps, so PyTorch's 1e-8 default applies. The Megatron backend uses adam_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 eps to AdamParams and the cross-rank config guard.

7. Truncated importance sampling falls back to the policy ratio

• Where: loss.py#L375-L385.
• What: with truncated_importance_sampling set and no original_logprobs, the loss is multiplied by clamp(prob_ratio, max=U), where prob_ratio is 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_logprobs are 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_logprobs is absent.

8. Unclamped exp(new - old) on the PPO path

• Where: loss.py#L299 and #L363-L367.
• What: if an old logprob is −inf, a sentinel such as −9999, or extremely stale, the ratio is inf. With ppo=True and 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_finite catches 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

• Where: tinker/service.py#L136-L180.
• What:
  • Masks are built as (group == g) | (parent == parent(g)). Prefix-tree roots have parent_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_async sits inside the per-row loop.
• Effect: wrong-context logprobs and gradients for multi-branch rows, plus extra optimizer steps per ART step.
• Confidence: plausible (logic, not numeric); not reproduced. A two-branch unit test would settle it.
• Suggested fix: build masks as (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 uses torch.std (n−1). The other backends use the population standard deviation (n) in preprocessing/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-L1020 and #L1414.
• What: (x @ A) @ B * s is 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 #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_grads is 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 before exp, 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.py core: 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).