diff --git a/blackwell/README.md b/blackwell/README.md index 1e4bd50104..3a51992bee 100644 --- a/blackwell/README.md +++ b/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`. diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index a06d81a512..46b6766edb 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index e7d9084ad8..3c0d5012ae 100644 --- a/unsloth/models/llama.py +++ b/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) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index a126a3dde1..2555f0df1f 100644 --- a/unsloth/models/rl_replacements.py +++ b/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)