training

Training forwards with a chunked prefill prefix and recurrent decode suffix.

Functions

get_linear_attention_layers	Return serving-policy layers from a module or iterable of model chunks.
linear_attention_training_phase	Supply per-sequence phases to a module or iterable of model chunks/layers.

get_linear_attention_layers(model, *, enabled_only=False, unphased_only=False)

Return serving-policy layers from a module or iterable of model chunks.

Disabled state quantizers are included by default so a baseline can use the same serving arithmetic and phase as QAT. Use enabled_only for setup checks and unphased_only to select layers needing adapter-provided defaults.

linear_attention_training_phase(model, prefill_lengths, *, sequence_lengths=None, cu_seqlens=None, preserve_existing=False)

Supply per-sequence phases to a module or iterable of model chunks/layers.

Enabled state quantization requires a phase, including calibration/eval. Disabled quantizers use the original module outside a phase and serving arithmetic inside. Megatron selective recompute captures the phase for backward after context exit.

sequence_lengths excludes right padding in dense rows or padded THD segments. cu_seqlens describes physical packed storage, including alignment padding. Megatron derives omitted valid lengths from unpadded cu_seqlens_q metadata; padded-only metadata requires explicit sequence_lengths. Padding produces zero attention outputs and never updates the recurrent state. preserve_existing=True installs defaults only on layers without a phase. Empty pipeline stages participate as no-ops.