Agent Lightning v1.0.3: Low-precision RL correction, reduced MoE routing memory

1. Stabilizing reward-based training with low-precision inference

The release adds score centering to the Qwen3.5-9B math training example. This correction removes average update bias caused by probability differences between the response-generating and training models. A low-precision response model and a model being trained may assign different probabilities to the next token. Score centering targets the resulting bias in training updates.

At each generation position, the example stores the top 128 candidate probabilities from token sampling. It compares them with the current training model's probabilities to calculate a correction. Probabilities for unrecorded candidates are approximated as proportional to those of the current model.

According to the company, score centering kept late-stage training rewards at about 0.6. The uncorrected GRPO baseline, which learns by comparing rewards across response groups, fell to 0 around step 45 and stayed there.

This result is a training curve showing performance changes in a math example using Qwen/Qwen3.5-9B on one B200. Response generation used INT8 weights and activations with an INT4 KV cache, while training ran in BF16. The reward uses math_verify to judge the final answer, assigning 1 for a correct answer and 0 otherwise. Truncated generations also receive 0.

2. The example separates inference and training precision

The release also includes model preparation and precision validation tools, plus a path to synchronize training weights with the INT8 inference model. Quantization, which reduces the bit width of numeric representations, targets linear layers that transform inputs into weighted sums. Embeddings, which represent tokens as numeric vectors, and the language model output layer remain in BF16. The settings belong in the inference section, as follows.

actor_rollout_ref:
  rollout:
    dtype: bfloat16
    quantization: int8
    quantization_config_file: ${oc.env:SC_INT8_CONFIG,outputs/precision_validation/Qwen3.5-9B-W8A8-INT8/config.json}
    engine_kwargs:
      vllm:
        linear_backend: triton
        attention_backend: TRITON_ATTN
        kv_cache_dtype: int4_per_token_head

dtype sets the precision of surrounding computations, while the explicit quantization and cache settings select the INT8 and INT4 paths. The full example trains for 200 updates with 32 prompts and 4 responses per prompt. A quick check runs 3 steps with 2 prompts and 4 responses per prompt.

3. MoE routing replay processes temporary arrays sequentially

MoE is a model architecture that selects a subset of experts to process tokens. R3 routing replay passes the expert selections made during response generation to training so it can replay the same routes. Previously, all routing arrays in a batch were decoded before being copied into the final tensor. The change decodes one array at a time and releases it after copying.

The contributor's published local CPU measurements cover 64 synthetic rollouts. As the table shows, both temporary space for decoded routing data and peak process resident memory, or RSS, decreased.

Table: Host memory for R3 routing replay

Conditions: Local CPU benchmark, 64 synthetic rollouts representing model-generated response batches

Unit: MB

The largest reduction is in temporary decoding buffers, while the final training tensor remains 128 MB before and after the change. The improvement reduces intermediate storage while preserving the output tensor's format and its padding and truncation behavior.

4. The cost and scope of the savings

Low-precision inference aims to improve memory and compute efficiency. However, score centering adds the work of recording and passing candidate probabilities and computing the training model's probabilities across the full vocabulary. Approximating the unrecorded distribution leaves room for correction errors. Both GRPO and score centering in this example normalize rewards by the group standard deviation, adjusting rewards by variation within the group. This is not an exact reproduction of the paper's mean-only setup and gradient descent, a training method that reduces loss by following gradients.

The release also preserves fractional reward values and improves normalization precision. It fixes an issue where converting rewards to BF16 first caused 1.0 and 1.001 to be treated as identical. It also prevents overflow in the normalization denominator, where values exceed the representable range, and avoids rounding token counts. These fixes reduce the loss of reinforcement learning update inputs during low-precision conversions or calculations with large numbers.

This version targets bottlenecks in low-precision inference and MoE routing replay separately, with memory savings coming from intermediate buffer handling rather than the final tensor.