monkeypatch.models.recurrent_support
monkeypatch.models.recurrent_support
Which recurrent / linear-attention architectures can train on packed sequences.
A packed row concatenates documents, and attention isolates them through the
mask or cu_seqlens. Recurrent mixers (Mamba, GatedDeltaNet, KDA, short
convolutions) carry state along the row instead, so each one needs a boundary
signal (seq_idx or cu_seqlens) threaded into its kernel, and each rank
under context parallelism needs the previous rank’s final state. Neither
failure raises: the loss just quietly trains on cross-document state.
Functions
| Name | Description |
|---|---|
| packs_async_grpo_rollouts | Async GRPO (and structured EBFT on it) packs every rank’s rollouts into one row. |
| recurrent_layer_types | The layer_types of a model config that are not attention. |
| trains_on_packed_rows | Whether batches concatenate samples, each starting where position_ids restart at 0. |
| validate_recurrent_model_config | Raise when a recurrent architecture would silently train on leaked state. |
packs_async_grpo_rollouts
monkeypatch.models.recurrent_support.packs_async_grpo_rollouts(cfg)Async GRPO (and structured EBFT on it) packs every rank’s rollouts into one row.
recurrent_layer_types
monkeypatch.models.recurrent_support.recurrent_layer_types(config)The layer_types of a model config that are not attention.
trains_on_packed_rows
monkeypatch.models.recurrent_support.trains_on_packed_rows(cfg)Whether batches concatenate samples, each starting where position_ids restart at 0.
validate_recurrent_model_config
monkeypatch.models.recurrent_support.validate_recurrent_model_config(
cfg,
model_config=None,
)Raise when a recurrent architecture would silently train on leaked state.