Accept instantiated models in AsyncGRPOTrainer

[kashif]

AsyncGRPOTrainer currently requires model to be a model ID and always loads the weights itself. This PR brings that constructor argument in line with the other trainers by also accepting an instantiated PreTrainedModel or torch.nn.Module.

This is the one missing primitive needed by external training backends such as Arctic Platform #84: an integration can pass a lightweight module carrying the model config, then keep its remote loss, optimizer, checkpointing, and weight-transfer behavior in its own subclass using the hooks that already exist.

Behavior

• A string model ID follows the existing loading path unchanged.
• An instantiated model is used directly.
• model_init_kwargs are ignored with a warning for an instantiated model, matching the established trainer behavior.
• When args is omitted, the default output name is derived from the instantiated model's config.

Verification

The existing end-to-end training test now passes a real preloaded tiny model, verifies that the trainer keeps that exact instance, trains it, and verifies that its parameters change. The existing model-ID construction test remains in place.

Scope

This PR intentionally does not add a training-client protocol, remote model/optimizer implementations, remote checkpoint handling, or remote-loss branches to TRL. It changes only AsyncGRPOTrainer and its test; external backends own those integration-specific pieces.

---

Note

Medium Risk
Constructor and model-loading paths diverge for string vs instance; wrong attention or ignored kwargs could affect training, but scope is limited to AsyncGRPOTrainer init with existing tree-attention checks unchanged for loaded models.

Overview
AsyncGRPOTrainer now accepts either a Hugging Face model id/path or an already-loaded PreTrainedModel, matching other TRL trainers.

For a string, behavior is unchanged: the trainer loads via create_model_from_path with tree attention, and model_init_kwargs is copied so the config dict is not mutated. For an instantiated model, that object is used as-is, attention is switched to TREE_ATTENTION for packed training, and model_init_kwargs is ignored with a warning. When args is omitted, the default run name comes from the model config (via get_config_model_id) instead of splitting a path string.

A new integration test test_train_model preloads a tiny Qwen model, checks identity and tree attention, runs train(), and asserts weights update—alongside the existing model-id test_train.

^{Reviewed by Cursor Bugbot for commit e7f6928. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[qgallouedec]

Thanks @kashif, and nice to see Snowflake-AI-Research/Arctic-Platform#84 running end to end on it.

I'd rather not merge this as is. AsyncGRPO is about to move out of experimental (@AmineDiro), and once it does, its constructor and protocols are public API we can't change freely. A third protocol plus GRPOLoss, _RemoteModel, _RemoteOptimizer and 11 if self._remote_model branches is a lot of surface to freeze for one backend, and every later change to the training loop would have to keep it working.

Most of it can (and should IMO) live on the Arctic side:

• optimizer: the constructor already takes optimizers=(...)
• save and resume: override save_model and _load_from_checkpoint
• weight sync: weight_transfer is already pluggable, and _streaming_iter is lazy, so a transfer that ignores the iterator gathers nothing
• forward and backward: override compute_loss

The one real blocker is that model must be a string, so the trainer always loads the full weights. If AsyncGRPOTrainer accepts an instantiated model like GRPOTrainer does (~10 lines), a placeholder goes through __init__ fine: the MoE checks, pad token and served model name only read model.config, and add_fused_lm_head only wraps forward.

So my proposal: shrink this PR / new PR to accepting an instantiated model, and keep the client in Arctic Platform. Rough sketch of your side (untested):

class RemoteModel(nn.Module):
    # Carries the config the trainer reads, and one parameter for Trainer/accelerate to place
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.generation_config = GenerationConfig.from_model_config(config)
        self.placeholder = nn.Parameter(torch.zeros(1))

    @property
    def dtype(self):  # read by the trainer's MFU setup
        return self.placeholder.dtype


class RemoteOptimizer(torch.optim.Optimizer):
    def __init__(self, params, client, lr):
        super().__init__(params, {"lr": lr})
        self.client = client

    def step(self, closure=None):
        # The trainer's scheduler writes the scheduled lr into param_groups
        self.client.optimizer_step(self.param_groups[0]["lr"])


class ArcticAsyncGRPOTrainer(AsyncGRPOTrainer):
    def __init__(self, *args, client, **kwargs):
        self.client = client
        super().__init__(*args, **kwargs)

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        input_ids, position_ids, completion_mask, old_log_probs, advantages = ...  # same unpacking as the parent
        # Forward on the backend, log probs come back detached
        log_probs = self.client.forward(input_ids, position_ids, completion_mask).requires_grad_()
        # On backward, ship d(loss)/d(log_probs) to the backend, which backprops sum(grad * log_probs)
        log_probs.register_hook(self.client.backward)
        loss = ...  # same GRPO loss as the parent, on log_probs
        return loss

    def save_model(self, output_dir=None, _internal_call=False):
        output_dir = output_dir or self.args.output_dir
        self.client.save(output_dir)
        self.processing_class.save_pretrained(output_dir)

    def _load_from_checkpoint(self, resume_from_checkpoint, model=None):
        self.client.load(resume_from_checkpoint)


client = ArcticClient(...)
model = RemoteModel(AutoConfig.from_pretrained("Qwen/Qwen3-8B"))
args = AsyncGRPOConfig(..., gradient_checkpointing=False)
trainer = ArcticAsyncGRPOTrainer(
    model=model,
    client=client,
    args=args,
    reward_funcs=...,
    train_dataset=...,
    optimizers=(RemoteOptimizer(model.parameters(), client, args.learning_rate), None),
    weight_transfer=ArcticWeightTransfer(client),  # syncs vLLM from the backend, ignores the iterator
)
trainer.train()

If the backend computes the loss next to the model instead, the same compute_loss override just sends old_log_probs and advantages along with the batch.

[qgallouedec]

Thks, that's the shape I had in mind. A few things to align with the other trainers:

• Non-string branch, like GRPO:

          else:
              model_revision = None
              if args.model_init_kwargs is not None:
                  logger.warning(
                      "You passed `model_init_kwargs` to the `AsyncGRPOConfig`, but your model is already instantiated. "
                      "The `model_init_kwargs` will be ignored."
                  )

  and drop the model_revision = None above the if.
• Type hint model: "str | PreTrainedModel", and in the docstring model (str or [~transformers.PreTrainedModel]) with a second bullet - A [~transformers.PreTrainedModel] object. (plus your sentence on the padding-free attention requirement). No other trainer lists torch.nn.Module.
• Tests: keep test_train on the model id and add a test_train_model that passes an instance, like tests/test_sft_trainer.py::test_train_model.

[qgallouedec]

otherwise lgtm, cc @AmineDiro

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 64ec7b2. Configure here.}

[AmineDiro]

With tree attention integrated this actually complicates the integration with some remotre training model, because the remote trainer will probably need to understand the shape of the data it gets OR its a completely different path and builds its own rollout worker API on to of the trainer.

[sfc-gh-truwase]

I reviewed the current head (e7f6928) against the remote-model integration that motivated accepting an instantiated nn.Module. There are two separate regressions.

1. Lightweight remote-model placeholders now fail in __init__

The non-string branch unconditionally calls:

model.set_attn_implementation(TREE_ATTENTION)

A remote backend does not instantiate the training model on the TRL process. Its local object is only a config/generation-config carrier plus one parameter for Trainer/Accelerate, and its subclass overrides compute_loss to run forward/backward remotely. Such a plain nn.Module has no set_attn_implementation, so construction now fails with AttributeError before the subclass can run.

Adding a no-op method to the placeholder or using hasattr would be unsafe: for a local model, this call is what prevents silently running a tree-packed batch with the wrong attention implementation.

I suggest moving all assumptions that TRL executes the local model into an explicit overridable hook, while keeping the default strict, e.g.:

def _prepare_model_for_compute_loss(self, model):
    if not isinstance(model, PreTrainedModel):
        raise TypeError("The default AsyncGRPO compute_loss requires a PreTrainedModel")
    model.set_attn_implementation(TREE_ATTENTION)
    self._validate_tree_attention_compatibility(model)
    add_fused_lm_head(model, temperature=self.temperature, outputs=("log_probs", "entropy"))
    return model

A remote-backend subclass can then deliberately return its config carrier unchanged. This is preferable to inferring intent from hasattr, and it gives the .sinks and VLM checks the same correct boundary: those checks are meaningful only for the model that actually executes the packed forward. (layer_types may remain a config-level validation if it is also a remote batching constraint.)

A smaller alternative is to gate the attention switch and weight-dependent checks on PreTrainedModel, but then a plain module using the stock compute_loss must still be rejected explicitly rather than allowed to train with incorrect attention.

2. The current tree-batch contract is independently incompatible with the remote loss path

Current AsyncGRPOTrainer.compute_loss consumes shift_labels, token_index, sampling_mask, segment_id, and tree_*. The remote integration still consumes token-aligned completion_mask, old_log_probs, and advantages; after bypassing the constructor failure its next result is KeyError("completion_mask").

This is not fixed by renaming fields. TreeRow.pack() deduplicates shared prefixes into trie nodes, and the backend needs tree_kv_index/path semantics to reproduce the same attention and output indexing. A conventional packed FA2 forward split by position_ids == 0 cannot reconstruct sibling paths and would silently compute incorrect log-probabilities and gradients.

Could this PR expose/document an injectable batching contract as well—either a non-tree collator for remote backends, or enough tree-batch inputs for the backend to implement the current semantics? The model-preparation hook fixes the constructor regression, but the remote integration cannot safely use current AsyncGRPOTrainer without resolving this second contract.

[qgallouedec]

Thanks @sfc-gh-truwase. Neither point needs a hook in TRL.

I suggest moving all assumptions that TRL executes the local model into an explicit overridable hook

The placeholder should be a PreTrainedModel, not a bare nn.Module. Snowflake-AI-Research/Arctic-Platform#84 already builds one (install_stub_model_loader: a 1-layer from_config model), so pass that instance directly: set_attn_implementation works on it and the trainer runs.

Adding a no-op method to the placeholder or using hasattr would be unsafe: for a local model, this call is what prevents silently running a tree-packed batch with the wrong attention implementation.

A no-op on your placeholder class leaves local models alone: they keep the real method. Only a hasattr in TRL would have that problem.

either a non-tree collator for remote backends, or enough tree-batch inputs for the backend to implement the current semantics?

Tree packing is the batch contract since #7328, and there is no flat path left. On your side, either run TRL's tree_attention on the workers (FA varlen plus an index_select on K/V, FA2 in #7762) and return per-position log-probs, or override get_train_dataloader to collate row.samples, which still carry completion_mask, old_log_probs and advantage. I'd rather not add a collator hook: its input would be list[TreeRow], an internal type we would then have to freeze.