From 3166cf8ab265b12a88753784933061ac2b7598be Mon Sep 17 00:00:00 2001 From: Varun Singh <23416001+sapiosaturn@users.noreply.github.com> Date: Sat, 21 Feb 2026 16:06:47 -0800 Subject: [PATCH 1/2] initial test for ring attn --- CHANGELOG.md | 1 + src/prime_rl/trainer/models/layers/attn.py | 71 +++++++++++++++++++ .../trainer/models/layers/ring_attn.py | 20 +++++- src/prime_rl/trainer/rl/train.py | 7 +- src/prime_rl/trainer/sft/train.py | 7 +- 5 files changed, 102 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3574ad8715..56cf3cf9be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -61,4 +61,5 @@ Documenting changes which affect configuration usage patterns (added/moved/remov - **`orchestrator.filters`**: Added orchestrator-side rollout filters for detecting degenerate generations. Supports `[[filters]] type = "gibberish"` (rare tokens at high entropy) and `[[filters]] type = "repetition"` (high-confidence token streaks). Detected rollouts get reward zeroed and completion mask cleared (2026-02-13) - **`inference.model.tool_call_parser`**: Changed default from `"hermes"` to auto-detection from model name. Uses `MODEL_TOOL_CALL_PARSER` dict to infer the correct vLLM tool call parser (e.g. Qwen3→`hermes`, GLM-4.5→`glm45`, GLM-4.7→`glm47`, MiniMax-M2→`minimax_m2`, INTELLECT-3→`hermes`). Unknown models default to `None`. Explicit values still take priority. (#1795, 2026-02-16) - **`orchestrator.eval.cancel_inflight_rollouts_on_eval`**: Added flag to optionally cancel in-flight training rollouts before starting online evals. When enabled, avoids congestion by preventing training and eval rollouts from running simultaneously, but slows training as the rollout pipeline must refill after each eval (default: False) (2026-02-16) +- **`model.cp` + AFMoE**: Context parallelism now works with AFMoE models. Added `substitute_afmoe_flash_attn` which patches `AfmoeFlashAttention` to use ring attention with per-layer `window_size` (sliding window for local layers, full attention for global layers). Also plumbed `window_size` through the FA3 ring attention wrapper (`ring_fa3_varlen_func`). Backward-compatible: all new parameters default to full attention `(-1, -1)`, existing models and configs are unaffected. (2026-02-21) diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index e326155c00..02f6b07790 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -323,3 +323,74 @@ def forward( return attn_output, attn_weights FlashAttention.forward = RingFlashAttention.forward + + +def substitute_afmoe_flash_attn( + process_group: torch.distributed.ProcessGroup, + heads_k_stride: int, + attn_impl: str = "flash_attention_2", +) -> None: + from ring_flash_attn import llama3_flash_attn_varlen_func + + from .ring_attn import ring_fa3_varlen_func + + use_fa3 = attn_impl == "flash_attention_3" + ring_func = ring_fa3_varlen_func if use_fa3 else llama3_flash_attn_varlen_func + + from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention + + class RingAfmoeFlashAttention(AfmoeFlashAttention): + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape) + key_states = self.k_proj(hidden_states).view(hidden_shape) + value_states = self.v_proj(hidden_states).view(hidden_shape) + gate_states = self.gate_proj(hidden_states) + + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if self.is_local_attention: + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + + from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS + + window_size = (self.sliding_window - 1, 0) if self.sliding_window is not None else (-1, -1) + + out = ring_func( + query_states[0], + key_states[0], + value_states[0], + cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"], + cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"], + max_seqlen_q=DATA_PARAMS["max_seqlen_q"], + max_seqlen_k=DATA_PARAMS["max_seqlen_k"], + local_k_slice=DATA_PARAMS["local_k_slice"], + causal=True, + window_size=window_size, + group=process_group, + heads_k_stride=heads_k_stride, + ) + if isinstance(out, tuple): + out = out[0] + + attn_output = out.contiguous().view(*input_shape, -1) + attn_output = attn_output * torch.sigmoid(gate_states) + attn_output = self.o_proj(attn_output) + return attn_output, None + + AfmoeFlashAttention.forward = RingAfmoeFlashAttention.forward diff --git a/src/prime_rl/trainer/models/layers/ring_attn.py b/src/prime_rl/trainer/models/layers/ring_attn.py index bfa34fa671..e76dddaa52 100644 --- a/src/prime_rl/trainer/models/layers/ring_attn.py +++ b/src/prime_rl/trainer/models/layers/ring_attn.py @@ -15,6 +15,7 @@ def _fa3_varlen_forward( max_seqlen_k: int, softmax_scale: float, causal: bool, + window_size: tuple[int, int] = (-1, -1), ) -> tuple[torch.Tensor, torch.Tensor]: from flash_attn_interface import _flash_attn_forward @@ -32,6 +33,8 @@ def _fa3_varlen_forward( "causal": causal, } ) + if "window_size" in params: + params["window_size"] = window_size out, lse, _, _ = _flash_attn_forward(**params) return out, lse @@ -52,6 +55,7 @@ def _fa3_varlen_backward( dv: torch.Tensor, softmax_scale: float, causal: bool, + window_size: tuple[int, int] = (-1, -1), ) -> None: from flash_attn_interface import _flash_attn_backward @@ -75,6 +79,8 @@ def _fa3_varlen_backward( "causal": causal, } ) + if "window_size" in params: + params["window_size"] = window_size _flash_attn_backward(**params) @@ -100,6 +106,8 @@ def forward( heads_k_stride: int, causal: bool, group_name: str, + window_size_left: int = -1, + window_size_right: int = -1, ) -> torch.Tensor: group = dist.group.WORLD for pg in dist.distributed_c10d._world.pg_map: @@ -108,6 +116,7 @@ def forward( break local_k_slice = slice(local_k_slice_start, local_k_slice_stop) + window_size = (window_size_left, window_size_right) softmax_scale = q.shape[-1] ** (-0.5) out_list = [] lse_list = [] @@ -146,6 +155,7 @@ def forward( max_seqlen_k=max_seqlen_k, softmax_scale=softmax_scale, causal=causal, + window_size=window_size, ) out_list.append(out_i) lse_list.append(lse_i) @@ -161,6 +171,7 @@ def forward( ctx.heads_k_stride = heads_k_stride ctx.causal = causal ctx.group_name = group_name + ctx.window_size = window_size return out @staticmethod @@ -234,6 +245,7 @@ def backward(ctx, dout: torch.Tensor): dv=dv_i, softmax_scale=ctx.softmax_scale, causal=causal, + window_size=ctx.window_size, ) if heads_k_stride != nheads_k: @@ -250,8 +262,9 @@ def backward(ctx, dout: torch.Tensor): dv[:, i : i + heads_k_stride] = dv_i # Grads for: q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, - # local_k_slice_start, local_k_slice_stop, heads_k_stride, causal, group_name - return dq, dk, dv, None, None, None, None, None, None, None, None, None + # local_k_slice_start, local_k_slice_stop, heads_k_stride, causal, group_name, + # window_size_left, window_size_right + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None def ring_fa3_varlen_func( @@ -266,6 +279,7 @@ def ring_fa3_varlen_func( causal: bool, heads_k_stride: int, group: dist.ProcessGroup, + window_size: tuple[int, int] = (-1, -1), ) -> torch.Tensor: return _RingFA3Varlen.apply( q, @@ -280,4 +294,6 @@ def ring_fa3_varlen_func( heads_k_stride, causal, group.group_name, + window_size[0], + window_size[1], ) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 7daee30fd5..ea2a5a7e4e 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -5,7 +5,7 @@ # Import environment before any other imports # ruff: noqa: I001 -from prime_rl.trainer.models.layers.attn import substitute_prime_rl_flash_attn +from prime_rl.trainer.models.layers.attn import substitute_afmoe_flash_attn, substitute_prime_rl_flash_attn from prime_rl.trainer.rl.broadcast import setup_weight_broadcast from prime_rl.utils.act_offloading import maybe_activation_offloading import torch @@ -177,6 +177,11 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: heads_k_stride=1, attn_impl=config.model.attn, ) + substitute_afmoe_flash_attn( + parallel_dims.world_mesh["cp"].get_group(), + heads_k_stride=1, + attn_impl=config.model.attn, + ) # Optionally, resume training from a checkpoint progress = Progress() diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index e5837d9736..25293fe066 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -8,7 +8,7 @@ # Import environment before any other imports # ruff: noqa: I001 -from prime_rl.trainer.models.layers.attn import substitute_prime_rl_flash_attn +from prime_rl.trainer.models.layers.attn import substitute_afmoe_flash_attn, substitute_prime_rl_flash_attn from prime_rl.utils.act_offloading import maybe_activation_offloading import torch from torch.profiler import profile, ProfilerActivity, record_function @@ -99,6 +99,11 @@ def train(config: SFTTrainerConfig): heads_k_stride=1, attn_impl=config.model.attn, ) + substitute_afmoe_flash_attn( + parallel_dims.world_mesh["cp"].get_group(), + heads_k_stride=1, + attn_impl=config.model.attn, + ) # Set up checkpoint manager logger.info(f"Initializing checkpoint managers ({config.ckpt})") From 198da6c375327946dc951cda55a3cd6bb7513319 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sun, 22 Feb 2026 21:59:28 +0000 Subject: [PATCH 2/2] tmp --- CHANGELOG.md | 2 +- .../trainer/models/afmoe/modeling_afmoe.py | 27 +-- src/prime_rl/trainer/models/layers/attn.py | 197 ++++-------------- src/prime_rl/trainer/rl/train.py | 9 +- src/prime_rl/trainer/sft/train.py | 9 +- 5 files changed, 63 insertions(+), 181 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 56cf3cf9be..ecdf598752 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -61,5 +61,5 @@ Documenting changes which affect configuration usage patterns (added/moved/remov - **`orchestrator.filters`**: Added orchestrator-side rollout filters for detecting degenerate generations. Supports `[[filters]] type = "gibberish"` (rare tokens at high entropy) and `[[filters]] type = "repetition"` (high-confidence token streaks). Detected rollouts get reward zeroed and completion mask cleared (2026-02-13) - **`inference.model.tool_call_parser`**: Changed default from `"hermes"` to auto-detection from model name. Uses `MODEL_TOOL_CALL_PARSER` dict to infer the correct vLLM tool call parser (e.g. Qwen3→`hermes`, GLM-4.5→`glm45`, GLM-4.7→`glm47`, MiniMax-M2→`minimax_m2`, INTELLECT-3→`hermes`). Unknown models default to `None`. Explicit values still take priority. (#1795, 2026-02-16) - **`orchestrator.eval.cancel_inflight_rollouts_on_eval`**: Added flag to optionally cancel in-flight training rollouts before starting online evals. When enabled, avoids congestion by preventing training and eval rollouts from running simultaneously, but slows training as the rollout pipeline must refill after each eval (default: False) (2026-02-16) -- **`model.cp` + AFMoE**: Context parallelism now works with AFMoE models. Added `substitute_afmoe_flash_attn` which patches `AfmoeFlashAttention` to use ring attention with per-layer `window_size` (sliding window for local layers, full attention for global layers). Also plumbed `window_size` through the FA3 ring attention wrapper (`ring_fa3_varlen_func`). Backward-compatible: all new parameters default to full attention `(-1, -1)`, existing models and configs are unaffected. (2026-02-21) +- **`model.cp` + AFMoE**: Context parallelism now works with AFMoE models via unified `substitute_ring_attn` which patches `_compute_attention` on both `FlashAttention` and `AfmoeFlashAttention` to use ring attention. Sliding window layers automatically get per-layer `window_size`; full attention layers default to `(-1, -1)`. Also plumbed `window_size` through the FA3 ring attention wrapper (`ring_fa3_varlen_func`). (2026-02-21) diff --git a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py index ff53b37f1b..5ebe3197db 100644 --- a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py +++ b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py @@ -189,6 +189,19 @@ def __init__(self, config: AfmoeAttentionConfig, flash_attn_version: int = 4): if self._flash_attn_version == 4: self._flash_attn_call = torch._dynamo.disable(self.func) + def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen): + """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" + args = [q, k, v, cu_seqlens, cu_seqlens] + if self._flash_attn_version != 4: + args.extend([max_seqlen, max_seqlen]) + kwargs: dict = {"causal": True} + if self.sliding_window is not None: + kwargs["window_size"] = (self.sliding_window - 1, 0) + out = self._flash_attn_call(*args, **kwargs) + if isinstance(out, tuple): + out = out[0] + return out + def forward( self, hidden_states: torch.Tensor, @@ -209,7 +222,6 @@ def forward( key_states = self.k_norm(key_states) if self.is_local_attention: - # apply_rotary_pos_emb expects [batch, heads, seq, dim] query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) cos, sin = position_embeddings @@ -217,18 +229,7 @@ def forward( query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) - # Flash attention varlen expects [total_tokens, heads, dim] - args = [query_states[0], key_states[0], value_states[0], cu_seqlens, cu_seqlens] - if self._flash_attn_version != 4: - args.extend([max_seqlen, max_seqlen]) - - kwargs: dict = {"causal": True} - if self.sliding_window is not None: - kwargs["window_size"] = (self.sliding_window - 1, 0) - - out = self._flash_attn_call(*args, **kwargs) - if isinstance(out, tuple): - out = out[0] + out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) attn_output = out.contiguous().view(*input_shape, -1) attn_output = attn_output * torch.sigmoid(gate_states) diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index 02f6b07790..0161413089 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -91,6 +91,20 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2): if self._flash_attn_version == 4: self._flash_attn_call = torch._dynamo.disable(self.func) + def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen): + """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" + args = [q, k, v, cu_seqlens, cu_seqlens] + if self._flash_attn_version != 4: + args.extend([max_seqlen, max_seqlen]) + kwargs: dict = {"causal": True} + sliding_window = getattr(self, "sliding_window", None) + if sliding_window is not None: + kwargs["window_size"] = (sliding_window - 1, 0) + out = self._flash_attn_call(*args, **kwargs) + if isinstance(out, tuple): + out = out[0] + return out + def forward( self, hidden_states: torch.Tensor, @@ -129,22 +143,7 @@ def forward( key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) - args = [ - query_states[0], - key_states[0], - value_states[0], - cu_seqlens, - cu_seqlens, - ] - if self._flash_attn_version != 4: - args.extend([max_seqlen, max_seqlen]) - - out = self._flash_attn_call( - *args, - causal=True, - ) - if isinstance(out, tuple): - out = out[0] + out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) out = out.contiguous() attn_output = out.view(1, out.shape[0], -1) @@ -242,11 +241,12 @@ def forward( } -def substitute_prime_rl_flash_attn( +def substitute_ring_attn( process_group: torch.distributed.ProcessGroup, heads_k_stride: int, attn_impl: str = "flash_attention_2", ) -> None: + """Patch _compute_attention on FlashAttention (and AfmoeFlashAttention) to use ring attention.""" from ring_flash_attn import llama3_flash_attn_varlen_func from .ring_attn import ring_fa3_varlen_func @@ -254,143 +254,34 @@ def substitute_prime_rl_flash_attn( use_fa3 = attn_impl == "flash_attention_3" ring_func = ring_fa3_varlen_func if use_fa3 else llama3_flash_attn_varlen_func - class RingFlashAttention(FlashAttention): - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - if self.use_qk_norm and self.qk_norm_type == "per_layer": - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - - query_states = query_states.view(hidden_shape) - key_states = key_states.view(hidden_shape) - value_states = value_states.view(hidden_shape) - - if self.use_qk_norm and self.qk_norm_type == "per_head": - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) - - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS - - cu_seqlens_q = DATA_PARAMS["cu_seqlens_q"] - cu_seqlens_k = DATA_PARAMS["cu_seqlens_k"] - max_seqlen_q = DATA_PARAMS["max_seqlen_q"] - max_seqlen_k = DATA_PARAMS["max_seqlen_k"] - local_k_slice = DATA_PARAMS["local_k_slice"] - - # TODO: Can we optimize the rotary application instead of double transpose? - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) - out = ring_func( - query_states[0], - key_states[0], - value_states[0], - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - local_k_slice=local_k_slice, - causal=True, - group=process_group, - heads_k_stride=heads_k_stride, - ) - if isinstance(out, tuple): - out = out[0] - out = out.contiguous() - attn_output = out.view(1, out.shape[0], -1) - attn_weights = None - - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - FlashAttention.forward = RingFlashAttention.forward - - -def substitute_afmoe_flash_attn( - process_group: torch.distributed.ProcessGroup, - heads_k_stride: int, - attn_impl: str = "flash_attention_2", -) -> None: - from ring_flash_attn import llama3_flash_attn_varlen_func - - from .ring_attn import ring_fa3_varlen_func + def _ring_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): + from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS + + window_size = (-1, -1) + sliding_window = getattr(self, "sliding_window", None) + if sliding_window is not None: + window_size = (sliding_window - 1, 0) + + out = ring_func( + q, + k, + v, + cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"], + cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"], + max_seqlen_q=DATA_PARAMS["max_seqlen_q"], + max_seqlen_k=DATA_PARAMS["max_seqlen_k"], + local_k_slice=DATA_PARAMS["local_k_slice"], + causal=True, + window_size=window_size, + group=process_group, + heads_k_stride=heads_k_stride, + ) + if isinstance(out, tuple): + out = out[0] + return out - use_fa3 = attn_impl == "flash_attention_3" - ring_func = ring_fa3_varlen_func if use_fa3 else llama3_flash_attn_varlen_func + FlashAttention._compute_attention = _ring_compute_attention from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention - class RingAfmoeFlashAttention(AfmoeFlashAttention): - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states = self.q_proj(hidden_states).view(hidden_shape) - key_states = self.k_proj(hidden_states).view(hidden_shape) - value_states = self.v_proj(hidden_states).view(hidden_shape) - gate_states = self.gate_proj(hidden_states) - - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - - if self.is_local_attention: - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - - from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS - - window_size = (self.sliding_window - 1, 0) if self.sliding_window is not None else (-1, -1) - - out = ring_func( - query_states[0], - key_states[0], - value_states[0], - cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"], - cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"], - max_seqlen_q=DATA_PARAMS["max_seqlen_q"], - max_seqlen_k=DATA_PARAMS["max_seqlen_k"], - local_k_slice=DATA_PARAMS["local_k_slice"], - causal=True, - window_size=window_size, - group=process_group, - heads_k_stride=heads_k_stride, - ) - if isinstance(out, tuple): - out = out[0] - - attn_output = out.contiguous().view(*input_shape, -1) - attn_output = attn_output * torch.sigmoid(gate_states) - attn_output = self.o_proj(attn_output) - return attn_output, None - - AfmoeFlashAttention.forward = RingAfmoeFlashAttention.forward + AfmoeFlashAttention._compute_attention = _ring_compute_attention diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index ea2a5a7e4e..2fedc98e59 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -5,7 +5,7 @@ # Import environment before any other imports # ruff: noqa: I001 -from prime_rl.trainer.models.layers.attn import substitute_afmoe_flash_attn, substitute_prime_rl_flash_attn +from prime_rl.trainer.models.layers.attn import substitute_ring_attn from prime_rl.trainer.rl.broadcast import setup_weight_broadcast from prime_rl.utils.act_offloading import maybe_activation_offloading import torch @@ -172,12 +172,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: if parallel_dims.cp_enabled: substitute_hf_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1) - substitute_prime_rl_flash_attn( - parallel_dims.world_mesh["cp"].get_group(), - heads_k_stride=1, - attn_impl=config.model.attn, - ) - substitute_afmoe_flash_attn( + substitute_ring_attn( parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1, attn_impl=config.model.attn, diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 25293fe066..c8b2583d59 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -8,7 +8,7 @@ # Import environment before any other imports # ruff: noqa: I001 -from prime_rl.trainer.models.layers.attn import substitute_afmoe_flash_attn, substitute_prime_rl_flash_attn +from prime_rl.trainer.models.layers.attn import substitute_ring_attn from prime_rl.utils.act_offloading import maybe_activation_offloading import torch from torch.profiler import profile, ProfilerActivity, record_function @@ -94,12 +94,7 @@ def train(config: SFTTrainerConfig): if parallel_dims.cp_enabled: assert config.data.seq_len % parallel_dims.cp == 0, "Sequence length must be divisible by CP degree" substitute_hf_flash_attn(parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1) - substitute_prime_rl_flash_attn( - parallel_dims.world_mesh["cp"].get_group(), - heads_k_stride=1, - attn_impl=config.model.attn, - ) - substitute_afmoe_flash_attn( + substitute_ring_attn( parallel_dims.world_mesh["cp"].get_group(), heads_k_stride=1, attn_impl=config.model.attn,