diff --git a/README.md b/README.md index 64f106f2c1..be781a6072 100644 --- a/README.md +++ b/README.md @@ -38,14 +38,16 @@ Notebooks are beginner friendly. Read our [guide](https://docs.unsloth.ai/get-st - See detailed documentation for Unsloth [here](https://docs.unsloth.ai/) ## ⚡ Quickstart -### Linux or WSL: -``` +### Linux or WSL +```bash pip install unsloth ``` ### Windows -For Windows, ```pip install unsloth``` works only if you have Pytorch installed. For more info, read our [Windows Guide](https://docs.unsloth.ai/get-started/installing-+-updating/windows-installation). +For Windows, `pip install unsloth` works only if you have Pytorch installed. For more info, read our [Windows Guide](https://docs.unsloth.ai/get-started/installing-+-updating/windows-installation). ### Docker Use our official [Unsloth Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Read our [Docker Guide](https://docs.unsloth.ai/get-started/install-and-update/docker). +### Blackwell +For RTX 50x, B200, 6000 GPUs, simply do `pip install unsloth`. Read our [Blackwell Guide](https://docs.unsloth.ai/basics/training-llms-with-blackwell-rtx-50-series-and-unsloth) for more details. ## 🦥 Unsloth.ai News - 📣 **Memory-efficient RL** We're introducing even better RL. Our new kernels & algos allows faster RL with 50% less VRAM & 10× more context. [Read blog](https://docs.unsloth.ai/new/memory-efficient-rl) @@ -145,8 +147,14 @@ For **advanced installation instructions** or if you see weird errors during ins 1. Install `torch` and `triton`. Go to https://pytorch.org to install it. For example `pip install torch torchvision torchaudio triton` 2. Confirm if CUDA is installed correctly. Try `nvcc`. If that fails, you need to install `cudatoolkit` or CUDA drivers. -3. Install `xformers` manually. You can try installing `vllm` and seeing if `vllm` succeeds. Check if `xformers` succeeded with `python -m xformers.info` Go to https://github.com/facebookresearch/xformers. Another option is to install `flash-attn` for Ampere GPUs. -4. Double check that your versions of Python, CUDA, CUDNN, `torch`, `triton`, and `xformers` are compatible with one another. The [PyTorch Compatibility Matrix](https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix) may be useful. +3. Install `xformers` manually via: +```bash +pip install ninja +pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers +``` +Check if `xformers` succeeded with `python -m xformers.info` Go to https://github.com/facebookresearch/xformers. Another option is to install `flash-attn` for Ampere GPUs and ignore `xformers` +5. You can try installing `vllm` and seeing if `vllm` succeeds. +6. Double check that your versions of Python, CUDA, CUDNN, `torch`, `triton`, and `xformers` are compatible with one another. The [PyTorch Compatibility Matrix](https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix) may be useful. 5. Finally, install `bitsandbytes` and check it with `python -m bitsandbytes` ### Conda Installation (Optional) @@ -216,18 +224,26 @@ Or, run the below manually in a Python REPL: try: import torch except: raise ImportError('Install torch via `pip install torch`') from packaging.version import Version as V -v = V(torch.__version__) +import re +v = V(re.match(r"[0-9\.]{3,}", torch.__version__).group(0)) cuda = str(torch.version.cuda) is_ampere = torch.cuda.get_device_capability()[0] >= 8 -if cuda != "12.1" and cuda != "11.8" and cuda != "12.4": raise RuntimeError(f"CUDA = {cuda} not supported!") +USE_ABI = torch._C._GLIBCXX_USE_CXX11_ABI +if cuda not in ("11.8", "12.1", "12.4", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!") if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!") elif v <= V('2.1.1'): x = 'cu{}{}-torch211' elif v <= V('2.1.2'): x = 'cu{}{}-torch212' elif v < V('2.3.0'): x = 'cu{}{}-torch220' elif v < V('2.4.0'): x = 'cu{}{}-torch230' elif v < V('2.5.0'): x = 'cu{}{}-torch240' -elif v < V('2.6.0'): x = 'cu{}{}-torch250' +elif v < V('2.5.1'): x = 'cu{}{}-torch250' +elif v <= V('2.5.1'): x = 'cu{}{}-torch251' +elif v < V('2.7.0'): x = 'cu{}{}-torch260' +elif v < V('2.7.9'): x = 'cu{}{}-torch270' +elif v < V('2.8.0'): x = 'cu{}{}-torch271' +elif v < V('2.8.9'): x = 'cu{}{}-torch280' else: raise RuntimeError(f"Torch = {v} too new!") +if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!") x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "") print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"') ``` diff --git a/blackwell/README.md b/blackwell/README.md deleted file mode 100644 index b16ea52d03..0000000000 --- a/blackwell/README.md +++ /dev/null @@ -1,207 +0,0 @@ -# Unsloth Blackwell Compatibility - -For RTX 5060, RTX 5070, RTX 5080, RTX 5090 GPUs and also B200, B40, GB100, GB102, GB20* and GPUs listed in https://en.wikipedia.org/wiki/Blackwell_(microarchitecture) - -## Overview - -`Blackwell` (`sm100+`) requires all dependent libraries to be compiled with `cuda 12.8`. - -The core libs for running unsloth which have dependencies on `CUDA` version are: -- `bitsandbytes` - already has wheels built with `CUDA 12.8` so `pip install` should work out of the box -- `triton` - requires `triton>=3.3.1` -- `torch` - requires installing with `pip install torch --extra-index-url https://download.pytorch.org/whl/cu128` -- `vllm` - vLLM 0.10.0 supports Blackwell now, but use CUDA 12.8: `uv pip install -U vllm --torch-backend=cu128` -- `xformers` - (Optional) as of 6/26, `xformers` wheels are not yet built with `sm100+` enabled as support was only recently [added](https://github.com/facebookresearch/xformers/commit/d9b3b6e2b38ca485c89507ef8ac1fbef2723cdfa) so will require a source build (see below). - -## Installation - -### Using uv - -The installation order is important, since we want the overwrite bundled dependencies with specific versions (namely, `xformers` and `triton`). - -1) I prefer to use `uv` over `pip` as it's faster and better for resolving dependencies, especially for libraries which depend on `torch` but for which a specific `CUDA` version is required per this scenario. - - Install `uv` - - ```bash - curl -LsSf https://astral.sh/uv/install.sh | sh && source $HOME/.local/bin/env - ``` - - Create a project dir and venv: - - ```bash - mkdir 'unsloth-blackwell' && cd 'unsloth-blackwell' - uv venv .venv --python=3.12 --seed - source .venv/bin/activate - ``` - -2) Install `vllm` - - ```bash - uv pip install -U vllm --torch-backend=cu128 - ``` - - Note that we have to specify `cu128`, otherwise `vllm` will install `torch==2.7.0` but with `cu126`. - -3) Install `unsloth` dependencies - - ```bash - uv pip install unsloth unsloth_zoo bitsandbytes - ``` - - If you notice weird resolving issues due to Xformers, you can also install Unsloth from source without Xformers: - - ```bash - uv pip install -qqq \ - "unsloth_zoo[base] @ git+https://github.com/unslothai/unsloth-zoo" \ - "unsloth[base] @ git+https://github.com/unslothai/unsloth" - ``` - -4) Download and build `xformers` (Optional) - - Xformers is optional, but it is definitely faster and uses less memory. We'll use PyTorch's native SDPA if you do not want Xformers. Building Xformers from source might be slow, so beware! - - ```bash - # First uninstall xformers installed by previous libraries - uv pip uninstall xformers - - # Clone and build - git clone --depth=1 https://github.com/facebookresearch/xformers --recursive - cd xformers - export TORCH_CUDA_ARCH_LIST="12.0" - python setup.py install - ``` - - Note that we have to explicitly set `TORCH_CUDA_ARCH_LIST=12.0`. - -5) `transformers` - Install any transformers version, but best to get the latest. - - ```bash - uv pip install -U transformers - ``` - - -### Using conda or mamba - -1) Install `conda/mamba` - - ```bash - curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh" - ``` - - Run the installation script - ```bash - bash Miniforge3-$(uname)-$(uname -m).sh - ``` - - Create a conda or mamba environment - ```bash - conda create --name unsloth-blackwell python==3.12 -y - ``` - - Activate newly created environment - ```bash - conda activate unsloth-blackwell - ``` - -2) Install `vllm` - - Make sure you are inside the activated conda/mamba environment. You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:` - - ```bash - pip install -U vllm --extra-index-url https://download.pytorch.org/whl/cu128 - ``` - - Note that we have to specify `cu128`, otherwise `vllm` will install `torch==2.7.0` but with `cu126`. - -3) Install `unsloth` dependencies - - Make sure you are inside the activated conda/mamba environment. You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:` - - ```bash - pip install unsloth unsloth_zoo bitsandbytes - ``` - -4) Download and build `xformers` (Optional) - - Xformers is optional, but it is definitely faster and uses less memory. We'll use PyTorch's native SDPA if you do not want Xformers. Building Xformers from source might be slow, so beware! - - You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:` - - ```bash - # First uninstall xformers installed by previous libraries - pip uninstall xformers - - # Clone and build - git clone --depth=1 https://github.com/facebookresearch/xformers --recursive - cd xformers - export TORCH_CUDA_ARCH_LIST="12.0" - python setup.py install - ``` - - Note that we have to explicitly set `TORCH_CUDA_ARCH_LIST=12.0`. - -5) Update `triton` - - Make sure you are inside the activated conda/mamba environment. You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:` - - ```bash - pip install -U triton>=3.3.1 - ``` - - `triton>=3.3.1` is required for `Blackwell` support. - -6) `Transformers` - Install any transformers version, but best to get the latest. - - ```bash - uv pip install -U transformers - ``` - - -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 (reminder Xformers is optional, but faster for training) 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`. - -Note, might need to downgrade `numpy<=2.2` after all the installs. - -## Test -Both `test_llama32_sft.py` and `test_qwen3_grpo.py` should run without issue if correct install. If not, check diff between your installed env and `blackwell.requirements.txt`. diff --git a/blackwell/blackwell.requirements.txt b/blackwell/blackwell.requirements.txt deleted file mode 100644 index 124e7eb119..0000000000 --- a/blackwell/blackwell.requirements.txt +++ /dev/null @@ -1,180 +0,0 @@ -Using Python 3.12.11 environment at: unsloth-bw/.venv -accelerate==1.8.1 -aiohappyeyeballs==2.6.1 -aiohttp==3.12.13 -aiosignal==1.3.2 -airportsdata==20250622 -annotated-types==0.7.0 -anyio==4.9.0 -astor==0.8.1 -asttokens==3.0.0 -attrs==25.3.0 -bitsandbytes==0.46.0 -blake3==1.0.5 -cachetools==6.1.0 -certifi==2025.6.15 -charset-normalizer==3.4.2 -click==8.2.1 -cloudpickle==3.1.1 -comm==0.2.2 -compressed-tensors==0.10.1 -cupy-cuda12x==13.4.1 -cut-cross-entropy==25.1.1 -datasets==3.6.0 -debugpy==1.8.14 -decorator==5.2.1 -depyf==0.18.0 -diffusers==0.34.0 -dill==0.3.8 -diskcache==5.6.3 -distro==1.9.0 -dnspython==2.7.0 -docstring-parser==0.16 -einops==0.8.1 -email-validator==2.2.0 -executing==2.2.0 -fastapi==0.115.14 -fastapi-cli==0.0.7 -fastrlock==0.8.3 -filelock==3.18.0 -frozenlist==1.7.0 -fsspec==2025.5.1 -gguf==0.17.1 -h11==0.16.0 -hf-transfer==0.1.9 -hf-xet==1.1.5 -httpcore==1.0.9 -httptools==0.6.4 -httpx==0.28.1 -huggingface-hub==0.33.1 -idna==3.10 -importlib-metadata==8.7.0 -interegular==0.3.3 -ipykernel==6.29.5 -ipython==9.3.0 -ipython-pygments-lexers==1.1.1 -jedi==0.19.2 -jinja2==3.1.6 -jiter==0.10.0 -jsonschema==4.24.0 -jsonschema-specifications==2025.4.1 -jupyter-client==8.6.3 -jupyter-core==5.8.1 -lark==1.2.2 -llguidance==0.7.30 -llvmlite==0.44.0 -lm-format-enforcer==0.10.11 -markdown-it-py==3.0.0 -markupsafe==3.0.2 -matplotlib-inline==0.1.7 -mdurl==0.1.2 -mistral-common==1.6.2 -mpmath==1.3.0 -msgpack==1.1.1 -msgspec==0.19.0 -multidict==6.5.1 -multiprocess==0.70.16 -nest-asyncio==1.6.0 -networkx==3.5 -ninja==1.11.1.4 -numba==0.61.2 -numpy==2.2.0 -nvidia-cublas-cu12==12.8.3.14 -nvidia-cuda-cupti-cu12==12.8.57 -nvidia-cuda-nvrtc-cu12==12.8.61 -nvidia-cuda-runtime-cu12==12.8.57 -nvidia-cudnn-cu12==9.7.1.26 -nvidia-cufft-cu12==11.3.3.41 -nvidia-cufile-cu12==1.13.0.11 -nvidia-curand-cu12==10.3.9.55 -nvidia-cusolver-cu12==11.7.2.55 -nvidia-cusparse-cu12==12.5.7.53 -nvidia-cusparselt-cu12==0.6.3 -nvidia-nccl-cu12==2.26.2 -nvidia-nvjitlink-cu12==12.8.61 -nvidia-nvtx-cu12==12.8.55 -openai==1.92.2 -opencv-python-headless==4.11.0.86 -outlines==0.1.11 -outlines-core==0.1.26 -packaging==25.0 -pandas==2.3.0 -parso==0.8.4 -partial-json-parser==0.2.1.1.post6 -peft==0.15.2 -pexpect==4.9.0 -pillow==11.2.1 -pip==25.1.1 -platformdirs==4.3.8 -prometheus-client==0.22.1 -prometheus-fastapi-instrumentator==7.1.0 -prompt-toolkit==3.0.51 -propcache==0.3.2 -protobuf==3.20.3 -psutil==7.0.0 -ptyprocess==0.7.0 -pure-eval==0.2.3 -py-cpuinfo==9.0.0 -pyarrow==20.0.0 -pybase64==1.4.1 -pycountry==24.6.1 -pydantic==2.11.7 -pydantic-core==2.33.2 -pygments==2.19.2 -python-dateutil==2.9.0.post0 -python-dotenv==1.1.1 -python-json-logger==3.3.0 -python-multipart==0.0.20 -pytz==2025.2 -pyyaml==6.0.2 -pyzmq==27.0.0 -ray==2.47.1 -referencing==0.36.2 -regex==2024.11.6 -requests==2.32.4 -rich==14.0.0 -rich-toolkit==0.14.7 -rpds-py==0.25.1 -safetensors==0.5.3 -scipy==1.16.0 -sentencepiece==0.2.0 -setuptools==80.9.0 -shellingham==1.5.4 -shtab==1.7.2 -six==1.17.0 -sniffio==1.3.1 -stack-data==0.6.3 -starlette==0.46.2 -sympy==1.14.0 -tiktoken==0.9.0 -tokenizers==0.21.2 -torch==2.7.0+cu128 -torchaudio==2.7.0+cu128 -torchvision==0.22.0+cu128 -tornado==6.5.1 -tqdm==4.67.1 -traitlets==5.14.3 -transformers==4.52.4 -triton==3.3.1 -trl==0.19.0 -typeguard==4.4.4 -typer==0.16.0 -typing-extensions==4.14.0 -typing-inspection==0.4.1 -tyro==0.9.24 -tzdata==2025.2 -unsloth==2025.6.8 -unsloth-zoo==2025.6.5 -urllib3==2.5.0 -uvicorn==0.34.3 -uvloop==0.21.0 -vllm==0.9.2.dev280+g04e1642e3 -watchfiles==1.1.0 -wcwidth==0.2.13 -websockets==15.0.1 -wheel==0.45.1 -xformers==0.0.32+ff490c3.d20250626 -xgrammar==0.1.19 -xxhash==3.5.0 -yarl==1.20.1 -zipp==3.23.0 diff --git a/blackwell/test_llama32_sft.py b/blackwell/test_llama32_sft.py deleted file mode 100644 index 6448343182..0000000000 --- a/blackwell/test_llama32_sft.py +++ /dev/null @@ -1,178 +0,0 @@ -from unsloth import FastLanguageModel -from transformers import ( - AutoModelForCausalLM, - DataCollatorForSeq2Seq, - AutoTokenizer, -) -from trl import SFTConfig, SFTTrainer -from unsloth.chat_templates import ( - get_chat_template, - standardize_sharegpt, - train_on_responses_only, -) -from datasets import load_dataset -from peft import AutoPeftModelForCausalLM -import torch - -max_seq_length = 2048 -dtype = None -load_in_4bit = True - -fourbit_models = [ - "unsloth/Meta-Llama-3.1-8B-bnb-4bit", - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", - "unsloth/Meta-Llama-3.1-70B-bnb-4bit", - "unsloth/Meta-Llama-3.1-405B-bnb-4bit", - "unsloth/Mistral-Small-Instruct-2409", - "unsloth/mistral-7b-instruct-v0.3-bnb-4bit", - "unsloth/Phi-3.5-mini-instruct", - "unsloth/Phi-3-medium-4k-instruct", - "unsloth/gemma-2-9b-bnb-4bit", - "unsloth/gemma-2-27b-bnb-4bit", - "unsloth/Llama-3.2-1B-bnb-4bit", - "unsloth/Llama-3.2-1B-Instruct-bnb-4bit", - "unsloth/Llama-3.2-3B-bnb-4bit", - "unsloth/Llama-3.2-3B-Instruct-bnb-4bit", - "unsloth/Llama-3.3-70B-Instruct-bnb-4bit", -] - -model, tokenizer = FastLanguageModel.from_pretrained( - model_name="unsloth/Llama-3.2-1B-Instruct", - max_seq_length=max_seq_length, - dtype=dtype, - load_in_4bit=load_in_4bit, -) - -model: AutoModelForCausalLM = FastLanguageModel.get_peft_model( - model, - r=16, - target_modules=[ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - lora_alpha=16, - lora_dropout=0, - bias="none", - use_gradient_checkpointing="unsloth", - random_state=3407, - use_rslora=False, - loftq_config=None, -) - -tokenizer = get_chat_template(tokenizer, chat_template="llama-3.1") - - -def formatting_prompts_func(examples): - convos = examples["conversations"] - texts = [ - tokenizer.apply_chat_template( - convo, tokenize=False, add_generation_prompt=False - ) - for convo in convos - ] - return {"text": texts} - - -dataset = load_dataset("mlabonne/FineTome-100k", split="train") -dataset = standardize_sharegpt(dataset) -dataset = dataset.map(formatting_prompts_func, batched=True) -dataset[5]["conversations"] -dataset[5]["text"] - -trainer = SFTTrainer( - model=model, - tokenizer=tokenizer, - train_dataset=dataset, - dataset_text_field="text", - max_seq_length=max_seq_length, - data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer), - dataset_num_proc=2, - packing=False, - args=SFTConfig( - per_device_train_batch_size=2, - gradient_accumulation_steps=4, - warmup_steps=5, - max_steps=10, - learning_rate=2e-4, - logging_steps=1, - optim="adamw_8bit", - weight_decay=0.01, - lr_scheduler_type="linear", - seed=3407, - output_dir="outputs", - report_to="none", - ), -) - -trainer = train_on_responses_only( - trainer, - instruction_part="<|start_header_id|>user<|end_header_id|>\n\n", - response_part="<|start_header_id|>assistant<|end_header_id|>\n\n", -) - -tokenizer.decode(trainer.train_dataset[5]["input_ids"]) -space = tokenizer(" ", add_special_tokens=False).input_ids[0] -tokenizer.decode( - [space if x == -100 else x for x in trainer.train_dataset[5]["labels"]] -) - -gpu_stats = torch.cuda.get_device_properties(0) -start_gpu_memory = round( - torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3 -) -max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) -print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.") -print(f"{start_gpu_memory} GB of memory reserved.") - -trainer_stats = trainer.train() - -used_memory = round( - torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3 -) -used_memory_for_lora = round(used_memory - start_gpu_memory, 3) -used_percentage = round(used_memory / max_memory * 100, 3) -lora_percentage = round(used_memory_for_lora / max_memory * 100, 3) -print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.") -print( - f"{round(trainer_stats.metrics['train_runtime'] / 60, 2)} minutes used for training." -) -print(f"Peak reserved memory = {used_memory} GB.") -print(f"Peak reserved memory for training = {used_memory_for_lora} GB.") -print(f"Peak reserved memory % of max memory = {used_percentage} %.") -print( - f"Peak reserved memory for training % of max memory = {lora_percentage} %." -) - -tokenizer = get_chat_template(tokenizer, chat_template="llama-3.1") -FastLanguageModel.for_inference(model) - -messages = [ - { - "role": "user", - "content": "Continue the fibonnaci sequence: 1, 1, 2, 3, 5, 8,", - }, -] -inputs = tokenizer.apply_chat_template( - messages, - tokenize=True, - add_generation_prompt=True, - return_tensors="pt", -).to("cuda") - - -model.gradient_checkpointing_disable() # This is required if using transformers >= 4.53.0 and `use_cache=True` - -outputs = model.generate( - input_ids=inputs, - max_new_tokens=64, - use_cache=True, - temperature=1.5, - min_p=0.1, -) -print(tokenizer.batch_decode(outputs)) - diff --git a/blackwell/test_qwen3_grpo.py b/blackwell/test_qwen3_grpo.py deleted file mode 100644 index 6232ceed17..0000000000 --- a/blackwell/test_qwen3_grpo.py +++ /dev/null @@ -1,430 +0,0 @@ -from unsloth import FastLanguageModel -import torch - -max_seq_length = 2048 -lora_rank = 32 - -model, tokenizer = FastLanguageModel.from_pretrained( - model_name="unsloth/Qwen3-0.6B-Base", - max_seq_length=max_seq_length, - load_in_4bit=False, - fast_inference=True, - max_lora_rank=lora_rank, - gpu_memory_utilization=0.7, -) - -model = FastLanguageModel.get_peft_model( - model, - r=lora_rank, - target_modules=[ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - lora_alpha=lora_rank * 2, - use_gradient_checkpointing="unsloth", - random_state=3407, -) - -reasoning_start = "" -reasoning_end = "" -solution_start = "" -solution_end = "" - -system_prompt = f"""You are given a problem. -Think about the problem and provide your working out. -Place it between {reasoning_start} and {reasoning_end}. -Then, provide your solution between {solution_start}{solution_end}""" -system_prompt - -chat_template = ( - "{% if messages[0]['role'] == 'system' %}" - "{{ messages[0]['content'] + eos_token }}" - "{% set loop_messages = messages[1:] %}" - "{% else %}" - "{{ '{system_prompt}' + eos_token }}" - "{% set loop_messages = messages %}" - "{% endif %}" - "{% for message in loop_messages %}" - "{% if message['role'] == 'user' %}" - "{{ message['content'] }}" - "{% elif message['role'] == 'assistant' %}" - "{{ message['content'] + eos_token }}" - "{% endif %}" - "{% endfor %}" - "{% if add_generation_prompt %}{{ '{reasoning_start}' }}" - "{% endif %}" -) - -chat_template = chat_template.replace( - "'{system_prompt}'", f"'{system_prompt}'" -).replace("'{reasoning_start}'", f"'{reasoning_start}'") -tokenizer.chat_template = chat_template - -tokenizer.apply_chat_template( - [ - {"role": "user", "content": "What is 1+1?"}, - { - "role": "assistant", - "content": f"{reasoning_start}I think it's 2.{reasoning_end}{solution_start}2{solution_end}", - }, - {"role": "user", "content": "What is 2+2?"}, - ], - tokenize=False, - add_generation_prompt=True, -) - -from datasets import load_dataset -import pandas as pd -import numpy as np - -dataset = load_dataset("unsloth/OpenMathReasoning-mini", split="cot") -dataset = dataset.to_pandas()[["expected_answer", "problem", "generated_solution"]] - -is_number = pd.to_numeric( - pd.Series(dataset["expected_answer"]), errors="coerce" -).notnull() -dataset = dataset.iloc[np.where(is_number)[0]] - -dataset - - -def format_dataset(x): - expected_answer = x["expected_answer"] - problem = x["problem"] - - thoughts = x["generated_solution"] - thoughts = thoughts.replace("", "").replace("", "") - - thoughts = thoughts.strip() - final_prompt = ( - reasoning_start - + thoughts - + reasoning_end - + solution_start - + expected_answer - + solution_end - ) - return [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": problem}, - {"role": "assistant", "content": final_prompt}, - ] - - -dataset["Messages"] = dataset.apply(format_dataset, axis=1) - -tokenizer.apply_chat_template(dataset["Messages"][0], tokenize=False) - -dataset["N"] = dataset["Messages"].apply( - lambda x: len(tokenizer.apply_chat_template(x)) -) - -dataset = dataset.loc[dataset["N"] <= max_seq_length / 2].copy() -dataset.shape - -from datasets import Dataset - -dataset["text"] = tokenizer.apply_chat_template( - dataset["Messages"].values.tolist(), tokenize=False -) -dataset = Dataset.from_pandas(dataset) -dataset - -from trl import SFTTrainer, SFTConfig - -trainer = SFTTrainer( - model=model, - tokenizer=tokenizer, - train_dataset=dataset, - args=SFTConfig( - dataset_text_field="text", - per_device_train_batch_size=1, - gradient_accumulation_steps=1, - warmup_steps=5, - num_train_epochs=2, - learning_rate=2e-4, - logging_steps=5, - optim="adamw_8bit", - weight_decay=0.01, - lr_scheduler_type="linear", - seed=3407, - report_to="none", - ), -) - -trainer.train() - -text = tokenizer.apply_chat_template( - dataset[0]["Messages"][:2], - tokenize=False, - add_generation_prompt=True, -) - -from transformers import TextStreamer - -_ = model.generate( - **tokenizer(text, return_tensors="pt").to("cuda"), - temperature=0, - max_new_tokens=1024, - streamer=TextStreamer(tokenizer, skip_prompt=False), -) - -del dataset -torch.cuda.empty_cache() -import gc - -gc.collect() - -from datasets import load_dataset - -dataset = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train") -dataset - -dataset[0]["prompt"] - -dataset[0]["solution"] - - -def extract_hash_answer(text): - return text - - -extract_hash_answer(dataset[0]["solution"]) - -dataset = dataset.map( - lambda x: { - "prompt": [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": x["prompt"]}, - ], - "answer": extract_hash_answer(x["solution"]), - } -) -dataset[0] - -import re - -solution_end_regex = ( - r"[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?" -) - -match_format = re.compile( - rf"{reasoning_end}.*?" - rf"{solution_start}(.+?){solution_end_regex}" - rf"[\s]{{0,}}$", - flags=re.MULTILINE | re.DOTALL, -) -match_format - -match_format.findall( - f"Let me think!\n2\n", -) - -match_format.findall( - f"Let me think! 2 \n\n", -) - - -def match_format_exactly(completions, **kwargs): - scores = [] - for completion in completions: - score = 0 - response = completion[0]["content"] - if match_format.search(response) is not None: - score += 3.0 - scores.append(score) - return scores - - -def match_format_approximately(completions, **kwargs): - scores = [] - for completion in completions: - score = 0 - response = completion[0]["content"] - - score += 0.5 if response.count(reasoning_end) == 1 else -1.0 - score += 0.5 if response.count(solution_start) == 1 else -1.0 - score += 0.5 if response.count(solution_end) == 1 else -1.0 - scores.append(score) - return scores - - -def check_answer(prompts, completions, answer, **kwargs): - question = prompts[0][-1]["content"] - responses = [completion[0]["content"] for completion in completions] - - extracted_responses = [ - guess.group(1) if (guess := match_format.search(r)) is not None else None - for r in responses - ] - - scores = [] - for guess, true_answer in zip(extracted_responses, answer): - score = 0 - if guess is None: - scores.append(-2.0) - continue - if guess == true_answer: - score += 5.0 - elif guess.strip() == true_answer.strip(): - score += 3.5 - else: - try: - ratio = float(guess) / float(true_answer) - if ratio >= 0.9 and ratio <= 1.1: - score += 2.0 - elif ratio >= 0.8 and ratio <= 1.2: - score += 1.5 - else: - score -= 2.5 - except: - score -= 4.5 - scores.append(score) - return scores - - -match_numbers = re.compile( - solution_start + r".*?[\s]{0,}([-]?[\d\.\,]{1,})", flags=re.MULTILINE | re.DOTALL -) -print(match_numbers.findall(" 0.34 ")) -print(match_numbers.findall(" 123,456 ")) -print(match_numbers.findall(" -0.234 ")) -print(match_numbers.findall("17")) - -global PRINTED_TIMES -PRINTED_TIMES = 0 -global PRINT_EVERY_STEPS -PRINT_EVERY_STEPS = 5 - - -def check_numbers(prompts, completions, answer, **kwargs): - question = prompts[0][-1]["content"] - responses = [completion[0]["content"] for completion in completions] - - extracted_responses = [ - guess.group(1) if (guess := match_numbers.search(r)) is not None else None - for r in responses - ] - - scores = [] - global PRINTED_TIMES - global PRINT_EVERY_STEPS - if PRINTED_TIMES % PRINT_EVERY_STEPS == 0: - print( - "*" * 20 + f"Question:\n{question}", - f"\nAnswer:\n{answer[0]}", - f"\nResponse:\n{responses[0]}", - f"\nExtracted:\n{extracted_responses[0]}", - ) - PRINTED_TIMES += 1 - - for guess, true_answer in zip(extracted_responses, answer): - if guess is None: - scores.append(-2.5) - continue - try: - true_answer = float(true_answer.strip()) - guess = float(guess.strip().replace(",", "")) - scores.append(3.5 if guess == true_answer else -1.5) - except: - scores.append(0) - continue - return scores - - -tokenized = dataset.map( - lambda x: { - "tokens": tokenizer.apply_chat_template( - x["prompt"], add_generation_prompt=True, tokenize=True - ) - }, - batched=True, -) -print(tokenizer.decode(tokenized[0]["tokens"])) -tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])}) - -import numpy as np - -maximum_length = int(np.quantile(tokenized["L"], 0.9)) -print("Max Length = ", maximum_length) - -dataset = dataset.select(np.where(np.array(tokenized["L"]) <= maximum_length)[0]) -del tokenized - -max_prompt_length = maximum_length + 1 -max_completion_length = max_seq_length - max_prompt_length - -from vllm import SamplingParams - -vllm_sampling_params = SamplingParams( - min_p=0.1, - top_p=1.0, - top_k=-1, - seed=3407, - stop=[tokenizer.eos_token], - include_stop_str_in_output=True, -) - -from trl import GRPOConfig, GRPOTrainer - -training_args = GRPOConfig( - vllm_sampling_params=vllm_sampling_params, - temperature=1.0, - learning_rate=5e-6, - weight_decay=0.01, - warmup_ratio=0.1, - lr_scheduler_type="linear", - optim="adamw_8bit", - logging_steps=1, - per_device_train_batch_size=1, - gradient_accumulation_steps=1, - num_generations=4, - max_prompt_length=max_prompt_length, - max_completion_length=max_completion_length, - max_steps=10, - save_steps=100, - report_to="none", - output_dir="outputs", -) - -trainer = GRPOTrainer( - model=model, - processing_class=tokenizer, - reward_funcs=[ - match_format_exactly, - match_format_approximately, - check_answer, - check_numbers, - ], - args=training_args, - train_dataset=dataset, -) -trainer.train() - -text = "What is the sqrt of 101?" - -from vllm import SamplingParams - -sampling_params = SamplingParams( - temperature=1.0, - top_k=50, - max_tokens=1024, -) - - -output = ( - model.fast_generate( - [text], - sampling_params=sampling_params, - lora_request=None, - )[0] - .outputs[0] - .text -) - -print(output)