lint

Author: bwasti

torchtitan/experiments/deterministic_vllm_rl/batch_invariant_backward.py

@@ -42,6 +42,7 @@
 # Custom autograd Functions for vLLM operations
 # ============================================================================
 
+
 class SiluAndMulFunction(Function):
     """
     Autograd function for vLLM's SiluAndMul activation.
@@ -86,7 +87,7 @@ def backward(ctx, grad_output):
 
         where d_silu(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x)))
         """
-        x, = ctx.saved_tensors
+        (x,) = ctx.saved_tensors
 
         # Split input into gate and up
         d = x.shape[-1] // 2
@@ -176,6 +177,7 @@ def backward(ctx, grad_output):
 # Backward operation implementations for autograd
 # ============================================================================
 
+
 def matmul_backward_impl(grad_output, self, other, output_mask):
     """
     Backward pass for matmul: y = matmul(a, b)
@@ -324,8 +326,8 @@ def rms_norm_backward_impl(grad_output, input, weight, eps):
 # Registration
 # ============================================================================
 
-_batch_invariant_backward_MODE = False
-_batch_invariant_backward_LIB = None
+_batch_invariant_backward_mode = False
+_batch_invariant_backward_lib = None
 
 
 def patch_batch_invariant_with_gradients():
@@ -335,28 +337,35 @@ def patch_batch_invariant_with_gradients():
     implementations by registering the backward operations. vLLM handles all the
     forward passes, we just add gradient support.
     """
-    global _batch_invariant_backward_MODE, _batch_invariant_backward_LIB
+    global _batch_invariant_backward_mode, _batch_invariant_backward_lib
 
-    if _batch_invariant_backward_MODE:
+    if _batch_invariant_backward_mode:
         return
 
     # Get vLLM's batch_invariant library (already created by init_batch_invariance)
     from vllm.model_executor.layers import batch_invariant as vllm_bi
 
-    if not hasattr(vllm_bi, '_batch_invariant_LIB') or vllm_bi._batch_invariant_LIB is None:
+    if (
+        not hasattr(vllm_bi, "_batch_invariant_LIB")
+        or vllm_bi._batch_invariant_LIB is None
+    ):
         raise RuntimeError(
             "vLLM's batch_invariant mode is not initialized. "
             "Call init_batch_invariance() first."
         )
 
     # Use vLLM's existing library - don't destroy it!
-    _batch_invariant_backward_LIB = vllm_bi._batch_invariant_LIB
+    _batch_invariant_backward_lib = vllm_bi._batch_invariant_LIB
 
     # Just add the backward operations - everything else is already handled by vLLM
-    _batch_invariant_backward_LIB.impl("aten::matmul_backward", matmul_backward_impl, "CUDA")
-    _batch_invariant_backward_LIB.impl("aten::linear_backward", linear_backward_impl, "CUDA")
+    _batch_invariant_backward_lib.impl(
+        "aten::matmul_backward", matmul_backward_impl, "CUDA"
+    )
+    _batch_invariant_backward_lib.impl(
+        "aten::linear_backward", linear_backward_impl, "CUDA"
+    )
 
-    _batch_invariant_backward_MODE = True
+    _batch_invariant_backward_mode = True
 
 
 def enable_batch_invariant_backward_mode():
@@ -366,25 +375,28 @@ def enable_batch_invariant_backward_mode():
 
 def disable_batch_invariant_backward_mode():
     """Disable batch invariant backward mode."""
-    global _batch_invariant_backward_MODE, _batch_invariant_backward_LIB
+    global _batch_invariant_backward_mode, _batch_invariant_backward_lib
 
-    if _batch_invariant_backward_LIB is not None:
-        _batch_invariant_backward_LIB._destroy()
+    if _batch_invariant_backward_lib is not None:
+        _batch_invariant_backward_lib._destroy()
 
-    _batch_invariant_backward_MODE = False
-    _batch_invariant_backward_LIB = None
+    _batch_invariant_backward_mode = False
+    _batch_invariant_backward_lib = None
 
 
 def is_batch_invariant_backward_mode_enabled():
     """Check if batch invariant backward mode is enabled."""
-    return _batch_invariant_backward_MODE
+    return _batch_invariant_backward_mode
 
 
 # ============================================================================
 # Public API for gradient-enabled vLLM operations
 # ============================================================================
 
-def rms_norm_with_gradients(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+
+def rms_norm_with_gradients(
+    input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
     """
     RMS normalization with gradient support.
 

torchtitan/experiments/deterministic_vllm_rl/models/attention.py

@@ -14,11 +14,13 @@
 
 class VLLMCompatibleFlashAttention(torch.nn.Module):
     """Wrapper around FlashAttention as used by VLLM"""
+
     def __init__(self) -> None:
         super().__init__()
         self.flash_attn_varlen_func = flash_attn_varlen_func
-        from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
         from vllm.attention.utils.fa_utils import get_flash_attn_version
+        from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
+
         self.vllm_is_batch_invariant = vllm_is_batch_invariant
         self.fa_version = get_flash_attn_version()
 
@@ -51,17 +53,29 @@ def forward(
         # Create cumulative sequence lengths
         # cu_seqlens: [0, seq_len, 2*seq_len, ..., batch_size*seq_len]
         cu_seqlens = torch.arange(
-            0, (batch_size + 1) * seq_len, seq_len,
-            dtype=torch.int32, device=q.device
+            0, (batch_size + 1) * seq_len, seq_len, dtype=torch.int32, device=q.device
         )
 
         # Wrap Flash Attention with manual backward pass
         class FlashAttnWithBackward(torch.autograd.Function):
             @staticmethod
-            def forward(ctx, q, k, v, cu_seqlens, seq_len, scale, num_splits, flash_fn, fa_version):
+            def forward(
+                ctx,
+                q,
+                k,
+                v,
+                cu_seqlens,
+                seq_len,
+                scale,
+                num_splits,
+                flash_fn,
+                fa_version,
+            ):
                 # Call flash attention for forward (fast)
                 output = flash_fn(
-                    q, k, v,
+                    q,
+                    k,
+                    v,
                     cu_seqlens_q=cu_seqlens,
                     cu_seqlens_k=cu_seqlens,
                     max_seqlen_q=seq_len,
@@ -94,7 +108,9 @@ def backward(ctx, grad_output):
                 k_batch = k.reshape(batch_size, seq_len, num_heads, head_dim)
                 v_batch = v.reshape(batch_size, seq_len, num_heads, head_dim)
                 out_batch = output.reshape(batch_size, seq_len, num_heads, head_dim)
-                grad_out_batch = grad_output.reshape(batch_size, seq_len, num_heads, head_dim)
+                grad_out_batch = grad_output.reshape(
+                    batch_size, seq_len, num_heads, head_dim
+                )
 
                 # Transpose to (batch, num_heads, seq_len, head_dim)
                 q_t = q_batch.transpose(1, 2)
@@ -108,11 +124,16 @@ def backward(ctx, grad_output):
                 scores = torch.matmul(q_t, k_t.transpose(-2, -1)) * scale
 
                 # Apply causal mask
-                causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool), diagonal=1)
-                scores = scores.masked_fill(causal_mask, float('-inf'))
+                causal_mask = torch.triu(
+                    torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool),
+                    diagonal=1,
+                )
+                scores = scores.masked_fill(causal_mask, float("-inf"))
 
                 # Softmax
-                attn_weights = torch.nn.functional.softmax(scores, dim=-1)  # (B, H, N, N)
+                attn_weights = torch.nn.functional.softmax(
+                    scores, dim=-1
+                )  # (B, H, N, N)
 
                 # Backward through attention
                 # out = attn_weights @ v
@@ -140,18 +161,29 @@ def backward(ctx, grad_output):
                 grad_k_t = torch.matmul(grad_scores.transpose(-2, -1), q_t)
 
                 # Transpose back and reshape to varlen format
-                grad_q = grad_q_t.transpose(1, 2).reshape(total_tokens, num_heads, head_dim)
-                grad_k = grad_k_t.transpose(1, 2).reshape(total_tokens, num_heads, head_dim)
-                grad_v = grad_v_t.transpose(1, 2).reshape(total_tokens, num_heads, head_dim)
+                grad_q = grad_q_t.transpose(1, 2).reshape(
+                    total_tokens, num_heads, head_dim
+                )
+                grad_k = grad_k_t.transpose(1, 2).reshape(
+                    total_tokens, num_heads, head_dim
+                )
+                grad_v = grad_v_t.transpose(1, 2).reshape(
+                    total_tokens, num_heads, head_dim
+                )
 
                 return grad_q, grad_k, grad_v, None, None, None, None, None, None
 
         # Call Flash Attention varlen with custom backward
         output_varlen = FlashAttnWithBackward.apply(
-            q_varlen, k_varlen, v_varlen,
-            cu_seqlens, seq_len, scale,
+            q_varlen,
+            k_varlen,
+            v_varlen,
+            cu_seqlens,
+            seq_len,
+            scale,
             1 if self.vllm_is_batch_invariant() else 0,
-            self.flash_attn_varlen_func, self.fa_version,
+            self.flash_attn_varlen_func,
+            self.fa_version,
         )
 
         # Convert back to batch format

torchtitan/experiments/deterministic_vllm_rl/models/qwen3/model_vllm_compat.py

@@ -12,17 +12,17 @@
 
 from torchtitan.components.tokenizer import BaseTokenizer
 
-# Import from main torchtitan
-from torchtitan.models.qwen3.model.args import Qwen3ModelArgs
-from torchtitan.protocols.model import AttentionMasksType
-from torchtitan.protocols.train_spec import ModelProtocol
-
 # Import gradient-enabled operations from experiment utilities
 from torchtitan.experiments.deterministic_vllm_rl.batch_invariant_backward import (
     rms_norm_with_gradients,
     silu_and_mul_with_gradients,
 )
 
+# Import from main torchtitan
+from torchtitan.models.qwen3.model.args import Qwen3ModelArgs
+from torchtitan.protocols.model import AttentionMasksType
+from torchtitan.protocols.train_spec import ModelProtocol
+
 # Import from local experiment's models
 from ..attention import VLLMCompatibleFlashAttention