Fireworks loss normalization depends on fwd/bwd chunking after #781

[signalrush]

Summary

PR #781 has the right ownership goal: rLLM should define loss normalization explicitly, and Fireworks should receive GradAccNormalization.NONE rather than infer a denominator.

The problem is the scope of the client-side denominator. The current custom-loss implementation normalizes each forward/backward chunk independently. When one optimizer batch is split across multiple accumulated forward/backward calls, that produces a sum of chunk means rather than a mean over the complete optimizer batch.

The required invariant is:

Changing fwd_bwd_group_size must not change the loss represented by one optimizer step.

Multiple forward/backward calls per optimizer step

The fully async trainer supports:

K = mini_batch_size // fwd_bwd_group_size

It schedules K forward/backward passes and performs one optimizer step after all nonempty passes:

• unified_trainer.py accumulation loop
• AsyncTrainingConfig field definitions

For example:

rllm.async_training.mini_batch_size: 8
rllm.async_training.fwd_bwd_group_size: 1

normally means eight training-client calls, one per nonempty pass, followed by one optimizer step. Builtin GRPO uses forward_backward; custom losses such as dppo_tv use forward_backward_custom.

Correct client-owned global token mean

For pass j, let:

• S_j be its summed token-gradient vector;
• N_j be its active-token count;
• N_global = sum_j N_j over every accepted datum contributing to this optimizer step.

The desired optimizer-batch token mean is:

G_expected = (sum_j S_j) / N_global

If normalization is fully owned by rLLM, it must know N_global before the first backward call and scale every pass using that same denominator:

pass_j contribution = S_j / N_global

Accumulation then gives:

sum_j (S_j / N_global) = (sum_j S_j) / N_global

The server can use:

GradAccNormalization.NONE

because every contribution is already globally scaled according to an rLLM-computed denominator.

Custom-loss behavior after #781

The custom path now calls:

build_custom_loss(..., server_normalized=False)

The closure only sees the current pass, so it divides by local N_j:

• per-pass client normalization
• build_custom_loss local divisor
• GradAccNormalization.NONE

The accumulated custom-loss gradient is therefore:

G_actual = sum_j (S_j / N_j)

If all pass sizes are equal, G_actual = K * G_expected; with eight passes it is exactly 8x. With unequal token counts, shorter passes are overweighted by N_global / N_j, generally changing the scale and potentially changing the direction. Dividing the final result by K would still be an unweighted mean of pass means, not a token mean, unless every N_j were identical.

Sequence-mean modes have the same issue if the global sequence denominator is replaced by one denominator per pass.

Builtin-loss behavior after #781

Builtin raw token-sum plus GradAccNormalization.NONE is valid and additive across passes. The issue is that #781 uses raw token-sum even when loss_agg_mode=token-mean or a sequence-mean mode was requested:

• builtin raw-token-sum behavior and warning

For requested token mean:

G_actual_builtin = sum_j S_j
G_expected        = (sum_j S_j) / N_global

This differs by the optimizer batch's token count. Adam can be approximately invariant to a stable uniform rescaling when epsilon and clipping are inactive, but variable token counts, gradient clipping, epsilon, weight decay, and existing optimizer moments can make the training behavior materially different.

Why this is not evidence of prior double normalization

Before #781, server_normalized=True meant that the custom closure returned a raw sum; it did not divide client-side:

• pre-#781 raw custom sum
• pre-#781 optimizer-window normalization

That older implementation delegated one global division to Fireworks. We should replace that delegation with an rLLM-computed global denominator, but it was not a client division followed by the same server division.

Separately, one builtin-GRPO optimizer step reported:

grad_norm_pre_norm = 3316.261962890625
global-normalized grad_norm = 0.00046834436943754554
ratio = approximately 7,080,820

This builtin observation independently confirms that the optimizer applied one global NUM_LOSS_TOKENS divisor. It is not an end-to-end custom-loss observation; the custom path's raw-sum behavior follows from the linked code.

grad_norm_rms is per-parameter-element RMS telemetry:

||g||_2 / sqrt(number of gradient elements)

It does not rescale the optimizer gradient.

Required implementation direction

rLLM should own both the denominator and its accumulation scope:

1. Collect or prepare all trajectory chunks belonging to one optimizer step before issuing the first backward call.
2. Apply filtering/transformation first, so counts reflect the datums and masks that will actually train.
3. Compute the optimizer-batch normalization metadata in rLLM:

token-sum              -> denominator 1
token-mean             -> total active tokens across all passes
seq-mean-token-sum     -> total trainable sequences across all passes
seq-mean-token-mean    -> per-sequence token mean, then total trainable sequences

4. Give every pass the same optimizer-batch denominator.
5. Keep GradAccNormalization.NONE; Fireworks must not infer token or sequence counts.

For custom losses, extend build_custom_loss to accept an explicit global normalization value rather than deriving it from the current pass.

For builtin losses, rLLM must apply the client-computed global scale to each pass's loss contribution before sending it. If a builtin kernel exposes only advantages/weights, scaling those is exact only for loss terms linear in that quantity. Kernels containing non-advantage terms need an explicit client-provided loss scale or must use the custom-loss path; silently leaving the requested aggregation mode unapplied is not sufficient.

Dropped/empty passes contribute neither numerator nor denominator. The count must come from the final action/loss mask, not from a server-defined notion such as nonzero-gradient tokens.

Regression tests

For both custom and representable builtin paths:

1. Construct one fixed synthetic optimizer batch with unequal pass token counts.
2. Run it as one forward/backward call.
3. Run the same batch split across K > 1 accumulated calls.
4. Compare gradients within numerical tolerance for token-mean and both sequence-mean modes.
5. Confirm token-sum remains an intentional raw sum.
6. Include an empty/dropped pass and verify it changes neither numerator nor denominator.

The result of one optimizer step should be invariant to fwd_bwd_group_size.

Related: #781

[signalrush]

Design correction: preserve streaming with an explicit client-owned scale

We do not need to wait for all forward/backward groups before starting backward.

rLLM can preserve the current streaming pipeline:

pass 1: send raw loss sum; rLLM records N_1
pass 2: send raw loss sum; rLLM records N_2
...
pass K: send raw loss sum; rLLM records N_K

optimizer boundary:
    N_global = sum_j N_j
    gradient_scale = 1 / N_global

The training service has already accumulated:

G_raw = sum_j S_j

Applying the explicit client-provided scale before clipping and Adam gives:

gradient_scale * G_raw
= (1 / N_global) * sum_j S_j
= G_expected

This retains all desired properties:

• forward/backward begins as soon as each group arrives;
• rLLM owns the exact action-token/sequence count and aggregation semantics;
• no server-inferred counting is used;
• only one optimizer update occurs after the configured accumulation window;
• builtin losses, including non-advantage terms, are scaled uniformly;
• the result is invariant to fwd_bwd_group_size.

The installed Fireworks 1.2.1 API currently cannot express this. optim_step accepts only:

GradAccNormalization.NUM_LOSS_TOKENS
GradAccNormalization.NUM_SEQUENCES
GradAccNormalization.NONE

AdamParams has no arbitrary gradient-scale/divisor field. We need an API such as:

optim_step(
    adam_params,
    grad_accumulation_normalization=GradAccNormalization.NONE,
    gradient_scale=1.0 / global_loss_tokens,
)

or equivalently:

optim_step(..., gradient_divisor=global_loss_tokens)

The server should not infer or choose the denominator. It should only apply the scalar supplied by rLLM, in this order:

accumulated raw gradient
-> multiply by client-provided gradient_scale
-> emit normalized grad-norm metrics
-> gradient clipping
-> Adam moment/update

This cannot be emulated by changing the learning rate: gradient clipping, Adam moments/epsilon, and weight decay make learning-rate scaling non-equivalent to scaling the gradient before the optimizer.

Until such an explicit scale is available, the alternatives are either server-owned counting (which does not meet the ownership requirement) or waiting until the global denominator is known before backward. The explicit optimizer-step scale is the correct streaming solution.

[signalrush]

Fireworks contract and the double-normalization hypothesis

The current Fireworks SDK/docs describe the intended RL contract as:

raw-sum RL/GRPO loss
-> optim_step(grad_accumulation_normalization="num_loss_tokens")
-> token-mean gradient

The installed SDK says:

Callers whose loss function returns a raw sum (e.g. RL/GRPO) should pass
GradAccNormalization.NUM_LOSS_TOKENS so the server divides by total loss tokens.

It warns about double normalization only when the loss function already returns a token/sequence mean and NUM_LOSS_TOKENS/NUM_SEQUENCES is additionally requested.

The public Fireworks loss-function documentation gives the same rule:

https://docs.fireworks.ai/fine-tuning/training-api/loss-functions

For the rLLM run under discussion, loss_fn=grpo resolves to the native builtin path:

forward_backward(..., builtin GRPO/PPO)  # no client-side loss division
optim_step(..., NUM_LOSS_TOKENS)         # one server-side division

Therefore, there is no client-visible second token division in the current rLLM path. The native server kernel source is not present in this repository, so an undocumented internal mean cannot be ruled out purely by inspecting rLLM; however, such an internal mean would contradict the published SDK contract.

Mathematically, for a denominator N independent of model parameters and applied before clipping/Adam:

grad((sum_i loss_i) / N) = (sum_i grad(loss_i)) / N

Thus loss-level and accumulated-gradient scalar normalization are equivalent when they use the same denominator and scope. The practical differences here are:

• who defines/counts N;
• whether N means all action-mask tokens or server-defined nonzero-gradient tokens;
• whether N spans the complete optimizer batch or only one pass;
• whether different loss terms need different reductions;
• whether scaling occurs before clipping and optimizer moments.

If rLLM is to own normalization, #781 must actually construct the requested global loss reduction. For builtin GRPO, raw token-sum plus NONE does not do that. For custom losses split across multiple passes, local pass means plus NONE do not do that either.

Definitive black-box check

Before changing the training objective, run the same deterministic builtin datum batch twice from identical trainer/optimizer state:

1. one copy of the datum batch;
2. K exact copies of that batch.

Measure raw pre-normalization gradient norm:

• raw-sum builtin kernel: duplicated norm scales by approximately K;
• internally mean-normalized builtin kernel: duplicated norm stays approximately constant.

Then repeat with NUM_LOSS_TOKENS:

• raw-sum builtin + one global division: normalized norm stays approximately constant;
• internally mean-normalized builtin + another division: normalized norm shrinks by approximately 1/K.

This directly tests the opaque native kernel and should be completed before treating double normalization as the cause of the small gradient norm.