Fix ref weights being broadcast on the first async update after resume

[ishzgu]

What does this PR do?

This PR fixes an incorrect first rollout weight update after resuming async non-colocate training with a reference model enabled.

The first update_weights() call can broadcast reference model weights to the SGLang rollout engines instead of the resumed actor weights. In the affected runs, this caused a sharp reward drop and spikes in KL divergence and grad_norm immediately after resume. The metrics recovered after approximately 2–3 training steps.

This is a small bug fix in the Megatron actor initialization path. It does not introduce a new feature or refactor existing abstractions.

Reproduction

The issue can be reproduced under the following conditions:

• Async non-colocate training
• A reference model enabled through KL loss, such as --use-kl-loss or a nonzero --kl-coef
• Resume from a checkpoint where the trained actor weights differ from the reference model weights

Reproduction sequence:

1. Start async training with a reference model and save an actor checkpoint after training has progressed.
2. Resume training from that actor checkpoint while using the original reference checkpoint.
3. Observe the first rollout after initialization.
4. The first rollout uses reference weights instead of the resumed actor weights, causing reward to drop and KL divergence and grad_norm to spike.
5. The metrics recover after the actor training path restores the actor weights and later weight updates broadcast the correct model.

This is generally not visible during a fresh run because the actor load path can fall back to the reference checkpoint, making the initial actor and reference weights identical.

Observed behavior

The following W&B metrics show the abnormal spike immediately after checkpoint resume. train/kl_loss and train/grad_norm spike at the same training step and recover shortly afterward.

These figures demonstrate the observed first-step regression. They support the reproduction but do not independently establish the root cause.

KL loss

Gradient norm

Root cause

Actor initialization first backs up the actor weights:

self._active_model_tag = "actor"
self.weights_backuper.backup("actor")

When a reference model is enabled, load_other_checkpoint("ref", ...) loads the reference weights into self.model. At the end of load_other_checkpoint(), it also updates the active model tag:

self.weights_backuper.backup(model_tag)
self._active_model_tag = model_tag

Previously, initialization restored the actor only inside the offload_train branch:

if self.args.offload_train:
    self._switch_model("actor")
    self.sleep()

In async non-colocate mode, offload_train is disabled, so initialization can return with the reference weights still loaded. train_async.py then immediately calls:

actor_model.update_weights()

The async distributed updater broadcasts weights from the currently loaded self.model, which is still the reference model at this point.

Fix

Restore the actor whenever initialization leaves a non-actor model active, and keep only sleep() conditional on offload_train:

if self._active_model_tag != "actor":
    self._switch_model("actor")

if self.args.offload_train:
    self.sleep()

This makes the post-initialization model state consistent and ensures that the first rollout weight update uses the resumed actor weights.

Verification and limitations

The issue and root cause were reproduced in the reference-model resume scenario described above. The diagnosis was verified using an alternative defensive workaround that called _switch_model("actor") at the beginning of update_weights(). With this workaround, reference weights were no longer broadcast during the first update, and the first-step regression disappeared.

The exact initialization-side patch in this PR has not yet been exercised in a complete end-to-end training run. The change is placed at initialization because that is where the inconsistent active-model state originates.

Only the reference-model path has been reproduced and checked. Initialization can also load teacher or old_actor model tags, and the generic active-tag check will restore the actor after those paths as well. Those configurations have not been tested as part of this PR.

The reported issue does not affect colocate/offloaded training because the existing offload_train path already restores the actor.