diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index c1fd16c..57c9c53 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -43,15 +43,20 @@ class Executor: ) task_ids = [t.task_id for t in tasks] + position_ids = ( + torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device) + .unsqueeze(0) + .expand(batch_sz, -1) + ) + input_mask = position_ids.unsqueeze(-1) >= torch.arange( + prompt_len, device=self.device + ) with torch.inference_mode(): self.model( input_ids, - position_ids=torch.arange( - start_pos, prompt_len, dtype=torch.long, device=self.device - ) - .unsqueeze(0) - .expand(batch_sz, -1), + input_mask=input_mask, + position_ids=position_ids, paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device), ) @@ -84,7 +89,10 @@ class Executor: position_ids = torch.tensor( [t.next_pos for t in tasks], dtype=torch.long, device=self.device ) - total_len = position_ids.max().item() + 1 + total_len = max(t.next_pos for t in tasks) + 1 + input_mask = position_ids[:, None, None] >= torch.arange( + total_len, device=self.device + ) task_ids = [t.task_id for t in tasks] @@ -122,6 +130,7 @@ class Executor: with torch.inference_mode(): outputs = self.model( input_ids.unsqueeze(1), + input_mask=input_mask, paged_cache=self.kv_cache.bind_tasks( task_ids, total_len, diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py index e94977a..7467395 100644 --- a/astrai/model/components/attention.py +++ b/astrai/model/components/attention.py @@ -76,9 +76,8 @@ class GQA(nn.Module): rotary_emb: Tensor, attn_mask: Tensor = None, paged_cache: Optional[CacheView] = None, + is_causal: bool = False, ) -> Tensor: - is_causal = attn_mask is None - q = self._split_heads(self.q_proj(x), self.n_heads) k = self._split_heads(self.k_proj(x), self.n_kv_heads) v = self._split_heads(self.v_proj(x), self.n_kv_heads) @@ -163,9 +162,9 @@ class MLA(nn.Module): rotary_emb: Tensor, attn_mask: Tensor = None, paged_cache: Optional[CacheView] = None, + is_causal: bool = False, ) -> Tensor: bsz, seq_len, _ = x.size() - is_causal = attn_mask is None q = self.q_proj(x) q = q.view(bsz, seq_len, self.n_heads, self.head_dim) diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py index 1c79d2a..e686fd0 100644 --- a/astrai/model/components/decoder_block.py +++ b/astrai/model/components/decoder_block.py @@ -26,12 +26,14 @@ class DecoderBlock(nn.Module): rotary_emb: Tensor, attention_mask: Optional[Tensor] = None, paged_cache: Optional[CacheView] = None, + is_causal: bool = False, ) -> Tensor: attn_output = self.attention( self.input_norm(x), rotary_emb, attention_mask, paged_cache, + is_causal, ) x = attn_output + x x = self.mlp(self.post_attention_norm(x)) + x diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py index 4b83652..cd2b5c2 100644 --- a/astrai/model/encoder.py +++ b/astrai/model/encoder.py @@ -59,10 +59,10 @@ class EmbeddingEncoder(AutoModel): x = self.embed_tokens(input_ids) rotary_emb = self.rotary_embedding(x, position_ids) - attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False) + attn_mask = process_attention_mask(input_mask) for layer in self.layers: - x = layer(x, rotary_emb, attn_mask, paged_cache=None) + x = layer(x, rotary_emb, attn_mask) hidden_states = self.norm(x) diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index 123d1ef..117d757 100644 --- a/astrai/model/transformer.py +++ b/astrai/model/transformer.py @@ -15,32 +15,15 @@ from astrai.model.components.rope import RotaryEmbedding def process_attention_mask( - input_tensor: Tensor, - position_ids: Optional[Tensor], - input_mask: Optional[Tensor] = None, - is_causal: bool = False, + input_mask: Optional[Tensor], ) -> Optional[Tensor]: - if position_ids is None: - return None - if input_mask is not None and input_mask.dim() > 2: - return input_mask - - device = input_tensor.device - B = input_tensor.size(0) - T = position_ids.max().item() + 1 - if input_mask is None: - if position_ids.min().item() == 0 and is_causal: - return None - attend = torch.ones(B, 1, T, dtype=torch.bool, device=device) - else: - attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1) - - if is_causal: - causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device) - attend = attend & causal - - return attend.unsqueeze(1) + return None + if input_mask.dim() == 2: + return input_mask[:, None, None, :] + if input_mask.dim() == 3: + return input_mask[:, None, :, :] + return input_mask @AutoModel.register("autoregressive_lm") @@ -119,10 +102,11 @@ class AutoRegressiveLM(AutoModel): x = self.embed_tokens(input_ids) rotary_emb = self.rotary_embedding(x, position_ids) - attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=True) + attn_mask = process_attention_mask(input_mask) + use_sdpa_causal_mask = attn_mask is None for layer in self.layers: - x = layer(x, rotary_emb, attn_mask, paged_cache) + x = layer(x, rotary_emb, attn_mask, paged_cache, use_sdpa_causal_mask) hidden_states = self.norm(x) logits = self.lm_head(hidden_states) diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index a7c7464..b2ab9c8 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -395,7 +395,9 @@ class GRPOStrategy(BaseStrategy): # response tokens. get_logprobs shifts the mask by one position, so # the first response token's logprob (predicted from the last prompt # token) is correctly included. - full_masks = torch.cat([torch.zeros_like(prompt_expanded), masks_flat], dim=-1) + full_masks = torch.cat( + [torch.zeros_like(prompt_expanded, dtype=torch.bool), masks_flat], dim=-1 + ) # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len). # Response token logprobs occupy the last ``response_len`` positions