Merge branch 'main' into nightly

blackwell/README.md

@@ -169,6 +169,40 @@ The installation order is important, since we want the overwrite bundled depende
 If you are using mamba as your package just replace conda with mamba for all commands shown above.
 
 
+## WSL-Specific Notes
+
+If you're using WSL (Windows Subsystem for Linux) and encounter issues during xformers compilation, follow these additional steps:
+
+1. **Increase WSL Memory Limit**
+   Create or edit the WSL configuration file:
+   ```bash
+   # Create or edit .wslconfig in your Windows user directory
+   # (typically C:\Users\YourUsername\.wslconfig)
+   
+   # Add these lines to the file
+   [wsl2]
+   memory=16GB  # Minimum 16GB recommended for xformers compilation
+   processors=4  # Adjust based on your CPU cores
+   swap=2GB
+   localhostForwarding=true
+   ```
+   After making these changes, restart WSL:
+   ```powershell
+   wsl --shutdown
+   ```
+
+2. **Install xformers**
+   Use the following command to install xformers with optimized compilation for WSL:
+   ```bash
+   # Set CUDA architecture for Blackwell GPUs
+   export TORCH_CUDA_ARCH_LIST="12.0"
+   
+   # Install xformers from source with optimized build flags
+   pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
+   ```
+   
+   The `--no-build-isolation` flag helps avoid potential build issues in WSL environments.
+
 ## Post Installation notes:
 
 After installation, your environment should look similar to `blackwell.requirements.txt`.

unsloth/kernels/rope_embedding.py

@@ -93,7 +93,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
         # [TODO] Changing blocksize to head_dim//2 seems to have
         # some concurrency / un-deterministic issues.
         BLOCK_SIZE, num_warps = calculate_settings(head_dim//2) # (head_dim//2)
-        
+
         # group_size = 4 # 4 or 8, too large group_size can hurt performance.
         div : int
         mod : int
@@ -155,6 +155,9 @@ pass
 def fast_rope_embedding(Q, K, cos, sin):
     Q = Fast_RoPE_Embedding.apply(Q.transpose(1, 2), cos, sin).transpose(1, 2)
     K = Fast_RoPE_Embedding.apply(K.transpose(1, 2), cos, sin).transpose(1, 2)
+    # synchronize before cat to avoid race condition
+    torch.cuda.current_stream(Q.device).synchronize()
+
     return Q, K
 pass
 
@@ -198,5 +201,6 @@ pass
 def inplace_rope_embedding(Q, K, cos, sin, position_ids):
     Q = Slow_RoPE_Embedding.apply(Q, cos, sin, position_ids)
     K = Slow_RoPE_Embedding.apply(K, cos, sin, position_ids)
+    torch.cuda.current_stream(Q.device).synchronize()
     return Q, K
 pass

unsloth/models/llama.py

@@ -499,8 +499,6 @@ def LlamaAttention_fast_forward(
     #     else inplace_rope_embedding(Q, K, cos, sin, position_ids)
     # )
     Q, K = fast_rope_embedding(Q, K, cos, sin)
-    # synchronize before cat to avoid race condition
-    torch.cuda.current_stream(Q.device).synchronize()
 
     if past_key_value is not None:
         K = torch.cat([past_key_value[0], K], dim = 2)

unsloth/models/rl_replacements.py

@@ -352,7 +352,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
     # Just copy over from _get_per_token_logps replacement function above. For now this returns None anyway
     def _get_per_token_logps_and_entropies(self, model, input_ids, attention_mask, logits_to_keep, batch_size = None, compute_entropy = False, *args, **kwargs):
         if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
-            return {"logps": None, "entropies": None} # Unsloth efficient GRPO
+            return None, None  # logps, entropies Unsloth efficient GRPO
         # Otherwise, calculate normally:
         if not hasattr(self, '_autocast_dtype'):
             self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16
@@ -373,7 +373,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                 entropies = entropy_from_logits(logits)
 
             # logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
-            return {"logps": logits, "entropies": entropies}
+            return logits, entropies  # logps, entropies
             # input_ids = input_ids[:, -logits_to_keep:]
             # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
             # See https://github.com/huggingface/trl/issues/2770
@@ -430,7 +430,7 @@ def grpo_trainer_compute_loss(function_name, function):
             lambda model, input_ids, attention_mask, logits_to_keep, batch_size=None, compute_entropy=False: \
             self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) \
             if hasattr(self, "_get_per_token_logps") else \
-            self._get_per_token_logps_and_entropies(model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy)['logps']
+            self._get_per_token_logps_and_entropies(model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy)[0]  # logps
 
         per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep)