diff --git a/blackwell/README.md b/blackwell/README.md new file mode 100644 index 0000000000..fdf5d2bc90 --- /dev/null +++ b/blackwell/README.md @@ -0,0 +1,87 @@ +## Unsloth Blackwell Compatibility + +### 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` - safest is to use the nightly build: `uv pip install -U vllm --torch-backend=cu128 --extra-index-url https://wheels.vllm.ai/nightly` +- `xformers` - 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 + +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 --extra-index-url https://wheels.vllm.ai/nightly + ``` + + 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 + ``` + +4) Download and build `xformers` + + ```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) Update `triton` + + ```bash + uv pip install -U triton>=3.3.1 + ``` + + `triton>=3.3.1` is required for `Blackwell` support. + +6) `transformers` + `transformers >= 4.53.0` breaks `unsloth` inference. Specifically, `transformers` with `gradient_checkpointing` enabled will automatically [switch off caching](https://github.com/huggingface/transformers/blob/67ddc82fbc7e52c6f42a395b4a6d278c55b77a39/src/transformers/modeling_layers.py#L52-L59). + + When using `unsloth` `FastLanguageModel` to `generate` directly after training with `use_cache=True`, this will result in mismatch between expected and actual outputs [here](https://github.com/unslothai/unsloth/blob/bfa6a3678e2fb8097c5ece41d095a8051f099db3/unsloth/models/llama.py#L939). + + Temporary solution is to switch off `gradient_checkpointing` (e.g., `model.disable_gradient_checkpointing()`) before generation if using `4.53.0` or stick with `4.52.4` for now: + + ```bash + uv pip install -U transformers==4.52.4 + ``` + +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 new file mode 100644 index 0000000000..124e7eb119 --- /dev/null +++ b/blackwell/blackwell.requirements.txt @@ -0,0 +1,180 @@ +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 new file mode 100644 index 0000000000..6448343182 --- /dev/null +++ b/blackwell/test_llama32_sft.py @@ -0,0 +1,178 @@ +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 new file mode 100644 index 0000000000..bf56ad338c --- /dev/null +++ b/blackwell/test_qwen3_grpo.py @@ -0,0 +1,429 @@ +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, +) +model.disable_gradient_checkpointing() +output = ( + model.fast_generate( + [text], + sampling_params=sampling_params, + lora_request=None, + )[0] + .outputs[0] + .text +) + +print(output)