diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..2db0c7cdd53688c21fa7eb8b7ba1a329bdd60e15 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +data/merged.jsonl filter=lfs diff=lfs merge=lfs -text +gemma4_e2b_out/checkpoint-1000/tokenizer.json filter=lfs diff=lfs merge=lfs -text +gemma4_gguf/tokenizer.json filter=lfs diff=lfs merge=lfs -text +lora_weights_final/tokenizer.json filter=lfs diff=lfs merge=lfs -text +unsloth_compiled_cache/__pycache__/UnslothDPOTrainer.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text +unsloth_compiled_cache/__pycache__/UnslothGRPOTrainer.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text +unsloth_compiled_cache/__pycache__/UnslothKTOTrainer.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text +unsloth_compiled_cache/__pycache__/UnslothRLOOTrainer.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text diff --git a/data/merged.jsonl b/data/merged.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..dab00a2832f7c5a273e61b5bf7d8d37256b8c3fa --- /dev/null +++ b/data/merged.jsonl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c41510b2b0f6b95e6ae376da19a095eb2d034c1455b2e3d609d9a8f194e3c27 +size 1141102581 diff --git a/gemma-4-e2b-it.F16.gguf b/gemma-4-e2b-it.F16.gguf new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/gemma4_e2b_out/README.md b/gemma4_e2b_out/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c7d2ffeba988e4877d71edf0e3c870d50100edfc --- /dev/null +++ b/gemma4_e2b_out/README.md @@ -0,0 +1,59 @@ +--- +base_model: unsloth/gemma-4-e2b-it-unsloth-bnb-4bit +library_name: transformers +model_name: gemma4_e2b_out +tags: +- generated_from_trainer +- trl +- unsloth +- sft +licence: license +--- + +# Model Card for gemma4_e2b_out + +This model is a fine-tuned version of [unsloth/gemma-4-e2b-it-unsloth-bnb-4bit](https://huggingface.co/unsloth/gemma-4-e2b-it-unsloth-bnb-4bit). +It has been trained using [TRL](https://github.com/huggingface/trl). + +## Quick start + +```python +from transformers import pipeline + +question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?" +generator = pipeline("text-generation", model="None", device="cuda") +output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0] +print(output["generated_text"]) +``` + +## Training procedure + + + + + +This model was trained with SFT. + +### Framework versions + +- TRL: 1.7.0 +- Transformers: 5.5.0 +- Pytorch: 2.10.0+cu128 +- Datasets: 5.0.0 +- Tokenizers: 0.22.2 + +## Citations + + + +Cite TRL as: + +```bibtex +@software{vonwerra2020trl, + title = {{TRL: Transformers Reinforcement Learning}}, + author = {von Werra, Leandro and Belkada, Younes and Tunstall, Lewis and Beeching, Edward and Thrush, Tristan and Lambert, Nathan and Huang, Shengyi and Rasul, Kashif and Gallouédec, Quentin}, + license = {Apache-2.0}, + url = {https://github.com/huggingface/trl}, + year = {2020} +} +``` \ No newline at end of file diff --git a/gemma4_e2b_out/checkpoint-1000/README.md b/gemma4_e2b_out/checkpoint-1000/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7a2628dee0e8593d629689f1d074cefe86c35da5 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/README.md @@ -0,0 +1,210 @@ +--- +base_model: unsloth/gemma-4-e2b-it-unsloth-bnb-4bit +library_name: peft +pipeline_tag: text-generation +tags: +- base_model:adapter:unsloth/gemma-4-e2b-it-unsloth-bnb-4bit +- lora +- sft +- transformers +- trl +- unsloth +--- + +# Model Card for Model ID + + + + + +## Model Details + +### Model Description + + + + + +- **Developed by:** [More Information Needed] +- **Funded by [optional]:** [More Information Needed] +- **Shared by [optional]:** [More Information Needed] +- **Model type:** [More Information Needed] +- **Language(s) (NLP):** [More Information Needed] +- **License:** [More Information Needed] +- **Finetuned from model [optional]:** [More Information Needed] + +### Model Sources [optional] + + + +- **Repository:** [More Information Needed] +- **Paper [optional]:** [More Information Needed] +- **Demo [optional]:** [More Information Needed] + +## Uses + + + +### Direct Use + + + +[More Information Needed] + +### Downstream Use [optional] + + + +[More Information Needed] + +### Out-of-Scope Use + + + +[More Information Needed] + +## Bias, Risks, and Limitations + + + +[More Information Needed] + +### Recommendations + + + +Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations. + +## How to Get Started with the Model + +Use the code below to get started with the model. + +[More Information Needed] + +## Training Details + +### Training Data + + + +[More Information Needed] + +### Training Procedure + + + +#### Preprocessing [optional] + +[More Information Needed] + + +#### Training Hyperparameters + +- **Training regime:** [More Information Needed] + +#### Speeds, Sizes, Times [optional] + + + +[More Information Needed] + +## Evaluation + + + +### Testing Data, Factors & Metrics + +#### Testing Data + + + +[More Information Needed] + +#### Factors + + + +[More Information Needed] + +#### Metrics + + + +[More Information Needed] + +### Results + +[More Information Needed] + +#### Summary + + + +## Model Examination [optional] + + + +[More Information Needed] + +## Environmental Impact + + + +Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700). + +- **Hardware Type:** [More Information Needed] +- **Hours used:** [More Information Needed] +- **Cloud Provider:** [More Information Needed] +- **Compute Region:** [More Information Needed] +- **Carbon Emitted:** [More Information Needed] + +## Technical Specifications [optional] + +### Model Architecture and Objective + +[More Information Needed] + +### Compute Infrastructure + +[More Information Needed] + +#### Hardware + +[More Information Needed] + +#### Software + +[More Information Needed] + +## Citation [optional] + + + +**BibTeX:** + +[More Information Needed] + +**APA:** + +[More Information Needed] + +## Glossary [optional] + + + +[More Information Needed] + +## More Information [optional] + +[More Information Needed] + +## Model Card Authors [optional] + +[More Information Needed] + +## Model Card Contact + +[More Information Needed] +### Framework versions + +- PEFT 0.19.1 \ No newline at end of file diff --git a/gemma4_e2b_out/checkpoint-1000/adapter_config.json b/gemma4_e2b_out/checkpoint-1000/adapter_config.json new file mode 100644 index 0000000000000000000000000000000000000000..e1e34b7fd837cc424e8c00434a330db773803dca --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/adapter_config.json @@ -0,0 +1,44 @@ +{ + "alora_invocation_tokens": null, + "alpha_pattern": {}, + "arrow_config": null, + "auto_mapping": { + "base_model_class": "Gemma4ForConditionalGeneration", + "parent_library": "transformers.models.gemma4.modeling_gemma4", + "unsloth_fixed": true + }, + "base_model_name_or_path": "unsloth/gemma-4-e2b-it-unsloth-bnb-4bit", + "bias": "none", + "corda_config": null, + "ensure_weight_tying": false, + "eva_config": null, + "exclude_modules": null, + "fan_in_fan_out": false, + "inference_mode": true, + "init_lora_weights": true, + "layer_replication": null, + "layers_pattern": null, + "layers_to_transform": null, + "loftq_config": {}, + "lora_alpha": 32, + "lora_bias": false, + "lora_dropout": 0, + "lora_ga_config": null, + "megatron_config": null, + "megatron_core": "megatron.core", + "modules_to_save": null, + "peft_type": "LORA", + "peft_version": "0.19.1", + "qalora_group_size": 16, + "r": 32, + "rank_pattern": {}, + "revision": null, + "target_modules": "(?:.*?(?:language|text).*?(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer).*?(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))|(?:\\bmodel\\.layers\\.[\\d]{1,}\\.(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer)\\.(?:(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)))", + "target_parameters": null, + "task_type": "CAUSAL_LM", + "trainable_token_indices": null, + "use_bdlora": null, + "use_dora": false, + "use_qalora": false, + "use_rslora": false +} \ No newline at end of file diff --git a/gemma4_e2b_out/checkpoint-1000/adapter_model.safetensors b/gemma4_e2b_out/checkpoint-1000/adapter_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a5b8be875ba039719502c80669f4e4c4267c7e94 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/adapter_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f14e7d5c4c6cfe03e5f99ffea455059061c7edc6b443e909b0c236b0bc8913cb +size 101424416 diff --git a/gemma4_e2b_out/checkpoint-1000/chat_template.jinja b/gemma4_e2b_out/checkpoint-1000/chat_template.jinja new file mode 100644 index 0000000000000000000000000000000000000000..582bb4df95a62b60d248387fac933343a5878f83 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/chat_template.jinja @@ -0,0 +1,70 @@ +{{ bos_token }}{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} +{%- set thinking = enable_thinking is defined and enable_thinking -%} +{%- set loop_messages = messages -%} +{%- if messages[0]['role'] in ['system', 'developer'] or thinking -%} + {{ '<|turn>system +' }} + {%- if thinking -%} + {{ '<|think|> +' }} + {%- endif -%} + {%- if messages[0]['role'] in ['system', 'developer'] -%} + {{ messages[0]['content'] | trim }} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {{ ' +' }} +{%- endif -%} +{%- for message in loop_messages -%} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%} + {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }} + {%- endif -%} + {%- if (message['role'] == 'assistant') -%} + {%- set role = "model" -%} + {%- else -%} + {%- set role = message['role'] -%} + {%- endif -%} + {{ '<|turn>' + role + ' +' }} + {%- if message['content'] is string -%} + {%- if role == "model" -%} + {{ strip_thinking(message['content']) }} + {%- else -%} + {{ message['content'] | trim }} + {%- endif -%} + {%- elif message['content'] is iterable -%} + {%- for item in message['content'] -%} + {%- if item['type'] == 'audio' -%} + {{ '<|audio|>' }} + {%- elif item['type'] == 'image' -%} + {{ '<|image|>' }} + {%- elif item['type'] == 'video' -%} + {{ '<|video|>' }} + {%- elif item['type'] == 'text' -%} + {%- if role == "model" -%} + {{ strip_thinking(item['text']) }} + {%- else -%} + {{ item['text'] | trim }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{ raise_exception("Invalid content type") }} + {%- endif -%} + {{ ' +' }} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{'<|turn>model +'}} +{%- endif -%} diff --git a/gemma4_e2b_out/checkpoint-1000/optimizer.pt b/gemma4_e2b_out/checkpoint-1000/optimizer.pt new file mode 100644 index 0000000000000000000000000000000000000000..cb1b0a84cb342cf270a55394a9aa39d366049500 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/optimizer.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a2575ed92e8989cf22e760c7aeb0346eb2d9dcbca70e4e57e5e7dd63ffb23173 +size 98628429 diff --git a/gemma4_e2b_out/checkpoint-1000/rng_state.pth b/gemma4_e2b_out/checkpoint-1000/rng_state.pth new file mode 100644 index 0000000000000000000000000000000000000000..3ef66339b9befa098183fd5d69faed6838e526b0 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/rng_state.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f1d565802a8e26c4e8a31328752b7a7fdc186d9401aa008e65697d0ad8c22e33 +size 14645 diff --git a/gemma4_e2b_out/checkpoint-1000/scaler.pt b/gemma4_e2b_out/checkpoint-1000/scaler.pt new file mode 100644 index 0000000000000000000000000000000000000000..0405b11307032f43fcf2212da19c52823cc01ac4 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/scaler.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:14ae2a2128444abab378aa06c09a61a84665f758fcc19fc46f5789b0bc1b5665 +size 1383 diff --git a/gemma4_e2b_out/checkpoint-1000/scheduler.pt b/gemma4_e2b_out/checkpoint-1000/scheduler.pt new file mode 100644 index 0000000000000000000000000000000000000000..8509d7f3b9ba4965ef8cea07f1e19f35c7850f22 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/scheduler.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ada16671df03ce2c4c7b2196578593f96e5a9638c91cc68f984858ec0e816498 +size 1465 diff --git a/gemma4_e2b_out/checkpoint-1000/tokenizer.json b/gemma4_e2b_out/checkpoint-1000/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..1ff9f3e3439a939b971f9919e821bf87e835a503 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f +size 32169626 diff --git a/gemma4_e2b_out/checkpoint-1000/tokenizer_config.json b/gemma4_e2b_out/checkpoint-1000/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..24b1522bb32b0eb05b631b1a9f112751b8cb0fab --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/tokenizer_config.json @@ -0,0 +1,289 @@ +{ + "audio_token": "<|audio|>", + "backend": "tokenizers", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "bos_token": "", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eos_token": "", + "eot_token": "", + "escape_token": "<|\"|>", + "etc_token": "", + "etd_token": "", + "etr_token": "", + "extra_special_tokens": [ + "<|video|>" + ], + "image_token": "<|image|>", + "is_local": false, + "mask_token": "", + "model_max_length": 131072, + "model_specific_special_tokens": { + "audio_token": "<|audio|>", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eot_token": "", + "escape_token": "<|\"|>", + "etc_token": "", + "etd_token": "", + "etr_token": "", + "image_token": "<|image|>", + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>" + }, + "pad_token": "", + "padding_side": "right", + "processor_class": "Gemma4Processor", + "response_schema": { + "properties": { + "content": { + "type": "string" + }, + "role": { + "const": "assistant" + }, + "thinking": { + "type": "string" + }, + "tool_calls": { + "items": { + "properties": { + "function": { + "properties": { + "arguments": { + "additionalProperties": {}, + "type": "object", + "x-parser": "gemma4-tool-call" + }, + "name": { + "type": "string" + } + }, + "type": "object", + "x-regex": "call\\:(?P\\w+)(?P\\{.*\\})" + }, + "type": { + "const": "function" + } + }, + "type": "object" + }, + "type": "array", + "x-regex-iterator": "<\\|tool_call>(.*?)" + } + }, + "type": "object", + "x-regex": "(\\<\\|channel\\>thought\\n(?P.*?)\\)?(?P\\<\\|tool_call\\>.*\\)?(?P(?:(?!\\)(?!\\<\\|tool_response\\>).)+)?(?:\\|\\<\\|tool_response\\>)?" + }, + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>", + "tokenizer_class": "GemmaTokenizer", + "unk_token": "", + "added_tokens_decoder": { + "0": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "1": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "2": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "3": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "4": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "46": { + "content": "<|tool>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "47": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "48": { + "content": "<|tool_call>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "49": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "50": { + "content": "<|tool_response>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "51": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "52": { + "content": "<|\"|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "98": { + "content": "<|think|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "100": { + "content": "<|channel>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "101": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "105": { + "content": "<|turn>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "106": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "255999": { + "content": "<|image>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "256000": { + "content": "<|audio>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258880": { + "content": "<|image|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258881": { + "content": "<|audio|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258882": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258883": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258884": { + "content": "<|video|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + } +} diff --git a/gemma4_e2b_out/checkpoint-1000/trainer_state.json b/gemma4_e2b_out/checkpoint-1000/trainer_state.json new file mode 100644 index 0000000000000000000000000000000000000000..ee47dfda3debeb318c4aacc2c96cb8b5080817f1 --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/trainer_state.json @@ -0,0 +1,734 @@ +{ + "best_global_step": null, + "best_metric": null, + "best_model_checkpoint": null, + "epoch": 0.024417639302632223, + "eval_steps": 500, + "global_step": 1000, + "is_hyper_param_search": false, + "is_local_process_zero": true, + "is_world_process_zero": true, + "log_history": [ + { + "epoch": 0.00024417639302632224, + "grad_norm": 0.48828125, + "learning_rate": 1.8e-05, + "loss": 1.1311161994934082, + "step": 10 + }, + { + "epoch": 0.0004883527860526445, + "grad_norm": 1.046875, + "learning_rate": 3.8e-05, + "loss": 1.2175333976745606, + "step": 20 + }, + { + "epoch": 0.0007325291790789666, + "grad_norm": 0.55859375, + "learning_rate": 5.8e-05, + "loss": 0.8916121482849121, + "step": 30 + }, + { + "epoch": 0.000976705572105289, + "grad_norm": 0.349609375, + "learning_rate": 7.800000000000001e-05, + "loss": 0.588713550567627, + "step": 40 + }, + { + "epoch": 0.0012208819651316112, + "grad_norm": 0.2431640625, + "learning_rate": 9.8e-05, + "loss": 0.5149059772491456, + "step": 50 + }, + { + "epoch": 0.0014650583581579332, + "grad_norm": 0.1484375, + "learning_rate": 0.000118, + "loss": 0.4876229286193848, + "step": 60 + }, + { + "epoch": 0.0017092347511842554, + "grad_norm": 0.10400390625, + "learning_rate": 0.000138, + "loss": 0.4492696762084961, + "step": 70 + }, + { + "epoch": 0.001953411144210578, + "grad_norm": 0.0966796875, + "learning_rate": 0.00015800000000000002, + "loss": 0.4119694232940674, + "step": 80 + }, + { + "epoch": 0.0021975875372369, + "grad_norm": 0.1025390625, + "learning_rate": 0.00017800000000000002, + "loss": 0.42351255416870115, + "step": 90 + }, + { + "epoch": 0.0024417639302632224, + "grad_norm": 0.73828125, + "learning_rate": 0.00019800000000000002, + "loss": 0.413313627243042, + "step": 100 + }, + { + "epoch": 0.0026859403232895444, + "grad_norm": 0.11376953125, + "learning_rate": 0.00019995065603657316, + "loss": 0.4356351852416992, + "step": 110 + }, + { + "epoch": 0.0029301167163158664, + "grad_norm": 0.1376953125, + "learning_rate": 0.000199780146829205, + "loss": 0.3644850492477417, + "step": 120 + }, + { + "epoch": 0.003174293109342189, + "grad_norm": 0.07421875, + "learning_rate": 0.00019948807088287883, + "loss": 0.392529559135437, + "step": 130 + }, + { + "epoch": 0.003418469502368511, + "grad_norm": 0.11279296875, + "learning_rate": 0.00019907478404714436, + "loss": 0.42130446434020996, + "step": 140 + }, + { + "epoch": 0.0036626458953948333, + "grad_norm": 0.12109375, + "learning_rate": 0.00019854078984834903, + "loss": 0.4250969409942627, + "step": 150 + }, + { + "epoch": 0.003906822288421156, + "grad_norm": 0.10693359375, + "learning_rate": 0.0001978867388761685, + "loss": 0.3913136005401611, + "step": 160 + }, + { + "epoch": 0.004150998681447478, + "grad_norm": 0.08984375, + "learning_rate": 0.00019711342799096361, + "loss": 0.40640673637390134, + "step": 170 + }, + { + "epoch": 0.0043951750744738, + "grad_norm": 0.1005859375, + "learning_rate": 0.00019622179935292855, + "loss": 0.3612894773483276, + "step": 180 + }, + { + "epoch": 0.004639351467500122, + "grad_norm": 0.10302734375, + "learning_rate": 0.00019521293927421388, + "loss": 0.3791630268096924, + "step": 190 + }, + { + "epoch": 0.004883527860526445, + "grad_norm": 0.05712890625, + "learning_rate": 0.00019408807689542257, + "loss": 0.3566859245300293, + "step": 200 + }, + { + "epoch": 0.005127704253552767, + "grad_norm": 0.0830078125, + "learning_rate": 0.00019284858268809137, + "loss": 0.3679236888885498, + "step": 210 + }, + { + "epoch": 0.005371880646579089, + "grad_norm": 0.10888671875, + "learning_rate": 0.0001914959667849825, + "loss": 0.38951241970062256, + "step": 220 + }, + { + "epoch": 0.005616057039605411, + "grad_norm": 0.087890625, + "learning_rate": 0.00019003187714021938, + "loss": 0.380203104019165, + "step": 230 + }, + { + "epoch": 0.005860233432631733, + "grad_norm": 0.11767578125, + "learning_rate": 0.0001884580975215084, + "loss": 0.3717670202255249, + "step": 240 + }, + { + "epoch": 0.006104409825658056, + "grad_norm": 0.07275390625, + "learning_rate": 0.00018677654533689287, + "loss": 0.3654949426651001, + "step": 250 + }, + { + "epoch": 0.006348586218684378, + "grad_norm": 0.11572265625, + "learning_rate": 0.00018498926929868642, + "loss": 0.32950897216796876, + "step": 260 + }, + { + "epoch": 0.0065927626117107, + "grad_norm": 0.08984375, + "learning_rate": 0.00018309844692743283, + "loss": 0.35263752937316895, + "step": 270 + }, + { + "epoch": 0.006836939004737022, + "grad_norm": 0.111328125, + "learning_rate": 0.00018110638189893267, + "loss": 0.37081763744354246, + "step": 280 + }, + { + "epoch": 0.007081115397763345, + "grad_norm": 0.0625, + "learning_rate": 0.00017901550123756906, + "loss": 0.34837267398834226, + "step": 290 + }, + { + "epoch": 0.007325291790789667, + "grad_norm": 0.078125, + "learning_rate": 0.00017682835235935236, + "loss": 0.37078888416290284, + "step": 300 + }, + { + "epoch": 0.007569468183815989, + "grad_norm": 0.08447265625, + "learning_rate": 0.00017454759996828623, + "loss": 0.3861753702163696, + "step": 310 + }, + { + "epoch": 0.007813644576842312, + "grad_norm": 0.10498046875, + "learning_rate": 0.00017217602280983623, + "loss": 0.32866339683532714, + "step": 320 + }, + { + "epoch": 0.008057820969868634, + "grad_norm": 0.07568359375, + "learning_rate": 0.00016971651028545648, + "loss": 0.4176007270812988, + "step": 330 + }, + { + "epoch": 0.008301997362894956, + "grad_norm": 0.07666015625, + "learning_rate": 0.00016717205893229903, + "loss": 0.32620184421539306, + "step": 340 + }, + { + "epoch": 0.008546173755921278, + "grad_norm": 0.0732421875, + "learning_rate": 0.00016454576877239507, + "loss": 0.36540043354034424, + "step": 350 + }, + { + "epoch": 0.0087903501489476, + "grad_norm": 0.057861328125, + "learning_rate": 0.0001618408395357554, + "loss": 0.3834134578704834, + "step": 360 + }, + { + "epoch": 0.009034526541973922, + "grad_norm": 0.0771484375, + "learning_rate": 0.00015906056676199255, + "loss": 0.3537560701370239, + "step": 370 + }, + { + "epoch": 0.009278702935000244, + "grad_norm": 0.1044921875, + "learning_rate": 0.00015620833778521307, + "loss": 0.39911091327667236, + "step": 380 + }, + { + "epoch": 0.009522879328026566, + "grad_norm": 0.0830078125, + "learning_rate": 0.000153287627607073, + "loss": 0.3524473667144775, + "step": 390 + }, + { + "epoch": 0.00976705572105289, + "grad_norm": 0.09619140625, + "learning_rate": 0.00015030199466302353, + "loss": 0.37404372692108157, + "step": 400 + }, + { + "epoch": 0.010011232114079211, + "grad_norm": 0.0859375, + "learning_rate": 0.00014725507648690543, + "loss": 0.3956392765045166, + "step": 410 + }, + { + "epoch": 0.010255408507105533, + "grad_norm": 0.06787109375, + "learning_rate": 0.00014415058527917452, + "loss": 0.375267767906189, + "step": 420 + }, + { + "epoch": 0.010499584900131856, + "grad_norm": 0.1181640625, + "learning_rate": 0.00014099230338415728, + "loss": 0.3407683610916138, + "step": 430 + }, + { + "epoch": 0.010743761293158178, + "grad_norm": 0.068359375, + "learning_rate": 0.00013778407868184672, + "loss": 0.3405567407608032, + "step": 440 + }, + { + "epoch": 0.0109879376861845, + "grad_norm": 0.0810546875, + "learning_rate": 0.00013452981989985348, + "loss": 0.3458467960357666, + "step": 450 + }, + { + "epoch": 0.011232114079210822, + "grad_norm": 0.0732421875, + "learning_rate": 0.00013123349185122327, + "loss": 0.37671115398406985, + "step": 460 + }, + { + "epoch": 0.011476290472237144, + "grad_norm": 0.07275390625, + "learning_rate": 0.00012789911060392294, + "loss": 0.3458314180374146, + "step": 470 + }, + { + "epoch": 0.011720466865263466, + "grad_norm": 0.091796875, + "learning_rate": 0.00012453073858788026, + "loss": 0.3307004690170288, + "step": 480 + }, + { + "epoch": 0.01196464325828979, + "grad_norm": 0.0791015625, + "learning_rate": 0.00012113247964553888, + "loss": 0.333436393737793, + "step": 490 + }, + { + "epoch": 0.012208819651316111, + "grad_norm": 0.09814453125, + "learning_rate": 0.00011770847403195834, + "loss": 0.3636301517486572, + "step": 500 + }, + { + "epoch": 0.012452996044342433, + "grad_norm": 0.181640625, + "learning_rate": 0.00011426289337055119, + "loss": 0.3728192329406738, + "step": 510 + }, + { + "epoch": 0.012697172437368755, + "grad_norm": 0.134765625, + "learning_rate": 0.0001107999355706023, + "loss": 0.3645843505859375, + "step": 520 + }, + { + "epoch": 0.012941348830395077, + "grad_norm": 0.1201171875, + "learning_rate": 0.00010732381971276318, + "loss": 0.3898338556289673, + "step": 530 + }, + { + "epoch": 0.0131855252234214, + "grad_norm": 0.09375, + "learning_rate": 0.00010383878090875201, + "loss": 0.36258883476257325, + "step": 540 + }, + { + "epoch": 0.013429701616447721, + "grad_norm": 0.08740234375, + "learning_rate": 0.00010034906514152238, + "loss": 0.35727477073669434, + "step": 550 + }, + { + "epoch": 0.013673878009474043, + "grad_norm": 0.0732421875, + "learning_rate": 9.685892409218717e-05, + "loss": 0.3593540906906128, + "step": 560 + }, + { + "epoch": 0.013918054402500366, + "grad_norm": 0.10595703125, + "learning_rate": 9.337260996000002e-05, + "loss": 0.3762841701507568, + "step": 570 + }, + { + "epoch": 0.01416223079552669, + "grad_norm": 0.08984375, + "learning_rate": 8.989437028170537e-05, + "loss": 0.336023736000061, + "step": 580 + }, + { + "epoch": 0.014406407188553011, + "grad_norm": 0.09814453125, + "learning_rate": 8.642844275656957e-05, + "loss": 0.33524041175842284, + "step": 590 + }, + { + "epoch": 0.014650583581579333, + "grad_norm": 0.09228515625, + "learning_rate": 8.297905008339677e-05, + "loss": 0.392057204246521, + "step": 600 + }, + { + "epoch": 0.014894759974605655, + "grad_norm": 0.111328125, + "learning_rate": 7.955039481582097e-05, + "loss": 0.3882176399230957, + "step": 610 + }, + { + "epoch": 0.015138936367631977, + "grad_norm": 0.06640625, + "learning_rate": 7.614665424214193e-05, + "loss": 0.38143367767333985, + "step": 620 + }, + { + "epoch": 0.0153831127606583, + "grad_norm": 0.0859375, + "learning_rate": 7.277197529594257e-05, + "loss": 0.352864146232605, + "step": 630 + }, + { + "epoch": 0.015627289153684623, + "grad_norm": 0.07666015625, + "learning_rate": 6.943046950368944e-05, + "loss": 0.35377163887023927, + "step": 640 + }, + { + "epoch": 0.015871465546710945, + "grad_norm": 0.076171875, + "learning_rate": 6.612620797547087e-05, + "loss": 0.36560447216033937, + "step": 650 + }, + { + "epoch": 0.016115641939737267, + "grad_norm": 0.1142578125, + "learning_rate": 6.286321644497655e-05, + "loss": 0.3602383375167847, + "step": 660 + }, + { + "epoch": 0.01635981833276359, + "grad_norm": 0.0654296875, + "learning_rate": 5.964547036476099e-05, + "loss": 0.3708587646484375, + "step": 670 + }, + { + "epoch": 0.01660399472578991, + "grad_norm": 0.07080078125, + "learning_rate": 5.647689006276726e-05, + "loss": 0.33884291648864745, + "step": 680 + }, + { + "epoch": 0.016848171118816233, + "grad_norm": 0.072265625, + "learning_rate": 5.33613359660109e-05, + "loss": 0.3408441781997681, + "step": 690 + }, + { + "epoch": 0.017092347511842555, + "grad_norm": 0.11572265625, + "learning_rate": 5.0302603897244474e-05, + "loss": 0.3734944581985474, + "step": 700 + }, + { + "epoch": 0.017336523904868877, + "grad_norm": 0.07080078125, + "learning_rate": 4.7304420450332244e-05, + "loss": 0.33618974685668945, + "step": 710 + }, + { + "epoch": 0.0175807002978952, + "grad_norm": 0.08349609375, + "learning_rate": 4.437043844996952e-05, + "loss": 0.3502551794052124, + "step": 720 + }, + { + "epoch": 0.01782487669092152, + "grad_norm": 0.076171875, + "learning_rate": 4.150423250127845e-05, + "loss": 0.3836493492126465, + "step": 730 + }, + { + "epoch": 0.018069053083947843, + "grad_norm": 0.1064453125, + "learning_rate": 3.8709294634702376e-05, + "loss": 0.33480191230773926, + "step": 740 + }, + { + "epoch": 0.018313229476974165, + "grad_norm": 0.0693359375, + "learning_rate": 3.5989030051504434e-05, + "loss": 0.3809062480926514, + "step": 750 + }, + { + "epoch": 0.018557405870000487, + "grad_norm": 0.1064453125, + "learning_rate": 3.334675297505476e-05, + "loss": 0.3911257266998291, + "step": 760 + }, + { + "epoch": 0.01880158226302681, + "grad_norm": 0.11083984375, + "learning_rate": 3.078568261295933e-05, + "loss": 0.38447113037109376, + "step": 770 + }, + { + "epoch": 0.01904575865605313, + "grad_norm": 0.10693359375, + "learning_rate": 2.8308939234951726e-05, + "loss": 0.3653350830078125, + "step": 780 + }, + { + "epoch": 0.019289935049079453, + "grad_norm": 0.078125, + "learning_rate": 2.5919540371325e-05, + "loss": 0.3945863485336304, + "step": 790 + }, + { + "epoch": 0.01953411144210578, + "grad_norm": 0.11376953125, + "learning_rate": 2.362039713653581e-05, + "loss": 0.39480888843536377, + "step": 800 + }, + { + "epoch": 0.0197782878351321, + "grad_norm": 0.0771484375, + "learning_rate": 2.1414310682459802e-05, + "loss": 0.289493727684021, + "step": 810 + }, + { + "epoch": 0.020022464228158423, + "grad_norm": 0.0849609375, + "learning_rate": 1.930396878561983e-05, + "loss": 0.3639736890792847, + "step": 820 + }, + { + "epoch": 0.020266640621184745, + "grad_norm": 0.1005859375, + "learning_rate": 1.7291942572543807e-05, + "loss": 0.36269190311431887, + "step": 830 + }, + { + "epoch": 0.020510817014211067, + "grad_norm": 0.10009765625, + "learning_rate": 1.538068338724361e-05, + "loss": 0.3807779312133789, + "step": 840 + }, + { + "epoch": 0.02075499340723739, + "grad_norm": 0.10205078125, + "learning_rate": 1.3572519804629536e-05, + "loss": 0.39589340686798097, + "step": 850 + }, + { + "epoch": 0.02099916980026371, + "grad_norm": 0.09521484375, + "learning_rate": 1.1869654793500784e-05, + "loss": 0.393012261390686, + "step": 860 + }, + { + "epoch": 0.021243346193290033, + "grad_norm": 0.07470703125, + "learning_rate": 1.0274163032567163e-05, + "loss": 0.3827983379364014, + "step": 870 + }, + { + "epoch": 0.021487522586316355, + "grad_norm": 0.08251953125, + "learning_rate": 8.787988382772705e-06, + "loss": 0.3293968915939331, + "step": 880 + }, + { + "epoch": 0.021731698979342677, + "grad_norm": 0.0693359375, + "learning_rate": 7.412941519000527e-06, + "loss": 0.3682440519332886, + "step": 890 + }, + { + "epoch": 0.021975875372369, + "grad_norm": 0.09228515625, + "learning_rate": 6.1506977240444074e-06, + "loss": 0.3646101474761963, + "step": 900 + }, + { + "epoch": 0.02222005176539532, + "grad_norm": 0.11669921875, + "learning_rate": 5.002794847534764e-06, + "loss": 0.38763275146484377, + "step": 910 + }, + { + "epoch": 0.022464228158421643, + "grad_norm": 0.1142578125, + "learning_rate": 3.970631432305694e-06, + "loss": 0.34759066104888914, + "step": 920 + }, + { + "epoch": 0.022708404551447965, + "grad_norm": 0.07177734375, + "learning_rate": 3.0554650104861136e-06, + "loss": 0.35361154079437257, + "step": 930 + }, + { + "epoch": 0.022952580944474287, + "grad_norm": 0.1259765625, + "learning_rate": 2.2584105713904125e-06, + "loss": 0.36288425922393797, + "step": 940 + }, + { + "epoch": 0.02319675733750061, + "grad_norm": 0.099609375, + "learning_rate": 1.580439203075812e-06, + "loss": 0.3696069002151489, + "step": 950 + }, + { + "epoch": 0.02344093373052693, + "grad_norm": 0.06982421875, + "learning_rate": 1.0223769092211012e-06, + "loss": 0.3502499103546143, + "step": 960 + }, + { + "epoch": 0.023685110123553253, + "grad_norm": 0.072265625, + "learning_rate": 5.849036027684606e-07, + "loss": 0.3617737293243408, + "step": 970 + }, + { + "epoch": 0.02392928651657958, + "grad_norm": 0.12109375, + "learning_rate": 2.685522775541904e-07, + "loss": 0.38356838226318357, + "step": 980 + }, + { + "epoch": 0.0241734629096059, + "grad_norm": 0.09130859375, + "learning_rate": 7.370835893788508e-08, + "loss": 0.38213505744934084, + "step": 990 + }, + { + "epoch": 0.024417639302632223, + "grad_norm": 0.07568359375, + "learning_rate": 6.092342209607083e-10, + "loss": 0.3387150287628174, + "step": 1000 + } + ], + "logging_steps": 10, + "max_steps": 1000, + "num_input_tokens_seen": 0, + "num_train_epochs": 1, + "save_steps": 500, + "stateful_callbacks": { + "TrainerControl": { + "args": { + "should_epoch_stop": false, + "should_evaluate": false, + "should_log": false, + "should_save": true, + "should_training_stop": true + }, + "attributes": {} + } + }, + "total_flos": 9.089382042631987e+16, + "train_batch_size": 2, + "trial_name": null, + "trial_params": null +} diff --git a/gemma4_e2b_out/checkpoint-1000/training_args.bin b/gemma4_e2b_out/checkpoint-1000/training_args.bin new file mode 100644 index 0000000000000000000000000000000000000000..84f0ab7c7090156058d6e8f3b27b6a815cd8ea5a --- /dev/null +++ b/gemma4_e2b_out/checkpoint-1000/training_args.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01c81c440c08fc4b233b9ce62ac6c2e3f9a0369823edad1a08fde166c4472b8d +size 5841 diff --git a/gemma4_gguf/chat_template.jinja b/gemma4_gguf/chat_template.jinja new file mode 100644 index 0000000000000000000000000000000000000000..582bb4df95a62b60d248387fac933343a5878f83 --- /dev/null +++ b/gemma4_gguf/chat_template.jinja @@ -0,0 +1,70 @@ +{{ bos_token }}{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} +{%- set thinking = enable_thinking is defined and enable_thinking -%} +{%- set loop_messages = messages -%} +{%- if messages[0]['role'] in ['system', 'developer'] or thinking -%} + {{ '<|turn>system +' }} + {%- if thinking -%} + {{ '<|think|> +' }} + {%- endif -%} + {%- if messages[0]['role'] in ['system', 'developer'] -%} + {{ messages[0]['content'] | trim }} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {{ ' +' }} +{%- endif -%} +{%- for message in loop_messages -%} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%} + {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }} + {%- endif -%} + {%- if (message['role'] == 'assistant') -%} + {%- set role = "model" -%} + {%- else -%} + {%- set role = message['role'] -%} + {%- endif -%} + {{ '<|turn>' + role + ' +' }} + {%- if message['content'] is string -%} + {%- if role == "model" -%} + {{ strip_thinking(message['content']) }} + {%- else -%} + {{ message['content'] | trim }} + {%- endif -%} + {%- elif message['content'] is iterable -%} + {%- for item in message['content'] -%} + {%- if item['type'] == 'audio' -%} + {{ '<|audio|>' }} + {%- elif item['type'] == 'image' -%} + {{ '<|image|>' }} + {%- elif item['type'] == 'video' -%} + {{ '<|video|>' }} + {%- elif item['type'] == 'text' -%} + {%- if role == "model" -%} + {{ strip_thinking(item['text']) }} + {%- else -%} + {{ item['text'] | trim }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{ raise_exception("Invalid content type") }} + {%- endif -%} + {{ ' +' }} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{'<|turn>model +'}} +{%- endif -%} diff --git a/gemma4_gguf/config.json b/gemma4_gguf/config.json new file mode 100644 index 0000000000000000000000000000000000000000..67231e183e6f1e59e126a7e9de18ee04e9123da3 --- /dev/null +++ b/gemma4_gguf/config.json @@ -0,0 +1,193 @@ +{ + "architectures": [ + "Gemma4ForConditionalGeneration" + ], + "audio_config": { + "_name_or_path": "", + "architectures": null, + "attention_chunk_size": 12, + "attention_context_left": 13, + "attention_context_right": 0, + "attention_invalid_logits_value": -1000000000.0, + "attention_logit_cap": 50.0, + "chunk_size_feed_forward": 0, + "conv_kernel_size": 5, + "torch_dtype": "float16", + "gradient_clipping": 10000000000.0, + "hidden_act": "silu", + "hidden_size": 1024, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "is_encoder_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "model_type": "gemma4_audio", + "num_attention_heads": 8, + "num_hidden_layers": 12, + "output_attentions": false, + "output_hidden_states": false, + "output_proj_dims": 1536, + "problem_type": null, + "residual_weight": 0.5, + "return_dict": true, + "rms_norm_eps": 1e-06, + "subsampling_conv_channels": [ + 128, + 32 + ], + "use_clipped_linears": true + }, + "audio_token_id": 258881, + "boa_token_id": 256000, + "boi_token_id": 255999, + "bos_token_id": 2, + "torch_dtype": "float16", + "eoa_token_id": 258883, + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "eos_token_id": 1, + "image_token_id": 258880, + "initializer_range": 0.02, + "model_name": "unsloth/gemma-4-e2b-it-unsloth-bnb-4bit", + "model_type": "gemma4", + "pad_token_id": 0, + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attention_k_eq_v": false, + "bos_token_id": 2, + "torch_dtype": "float16", + "enable_moe_block": false, + "eos_token_id": 1, + "expert_intermediate_size": null, + "final_logit_softcapping": 30.0, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 1536, + "hidden_size_per_layer_input": 256, + "initializer_range": 0.02, + "intermediate_size": 6144, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "max_position_embeddings": 131072, + "model_type": "gemma4_text", + "moe_intermediate_size": null, + "num_attention_heads": 8, + "num_experts": null, + "num_global_key_value_heads": null, + "num_hidden_layers": 35, + "num_key_value_heads": 1, + "num_kv_shared_layers": 20, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": { + "partial_rotary_factor": 0.25, + "rope_theta": 1000000.0, + "rope_type": "proportional" + }, + "sliding_attention": { + "rope_theta": 10000.0, + "rope_type": "default" + } + }, + "sliding_window": 512, + "tie_word_embeddings": true, + "top_k_experts": null, + "use_bidirectional_attention": null, + "use_cache": true, + "use_double_wide_mlp": true, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144 + }, + "tie_word_embeddings": true, + "unsloth_fixed": true, + "unsloth_version": "2026.6.9", + "use_cache": false, + "video_token_id": 258884, + "vision_config": { + "_name_or_path": "", + "architectures": null, + "attention_bias": false, + "attention_dropout": 0.0, + "chunk_size_feed_forward": 0, + "default_output_length": 280, + "torch_dtype": "float16", + "global_head_dim": 64, + "head_dim": 64, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 768, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "intermediate_size": 3072, + "is_encoder_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "max_position_embeddings": 131072, + "model_type": "gemma4_vision", + "num_attention_heads": 12, + "num_hidden_layers": 16, + "num_key_value_heads": 12, + "output_attentions": false, + "output_hidden_states": false, + "patch_size": 16, + "pooling_kernel_size": 3, + "position_embedding_size": 10240, + "problem_type": null, + "return_dict": true, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 100.0, + "rope_type": "default" + }, + "standardize": false, + "use_clipped_linears": true + }, + "vision_soft_tokens_per_image": 280 +} \ No newline at end of file diff --git a/gemma4_gguf/generation_config.json b/gemma4_gguf/generation_config.json new file mode 100644 index 0000000000000000000000000000000000000000..edda3c1950276d89fa83f5c515fdc0a38734fc59 --- /dev/null +++ b/gemma4_gguf/generation_config.json @@ -0,0 +1,14 @@ +{ + "bos_token_id": 2, + "do_sample": true, + "eos_token_id": [ + 1, + 106, + 50 + ], + "pad_token_id": 0, + "temperature": 1.0, + "top_k": 64, + "top_p": 0.95, + "transformers_version": "5.5.0" +} diff --git a/gemma4_gguf/model.safetensors b/gemma4_gguf/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5cbd9b4a9fef6d884aa54647a4d86c147dd985cb --- /dev/null +++ b/gemma4_gguf/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:440203cafd5bac689bae28d22382b2bc67d5a39849b590cd88e862c191185d5c +size 10246621918 diff --git a/gemma4_gguf/temp_split_0de139d7_000.safetensors b/gemma4_gguf/temp_split_0de139d7_000.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..37f58a994dd947f7716a9d041ed0eb99334cbd35 --- /dev/null +++ b/gemma4_gguf/temp_split_0de139d7_000.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9ebcea1985a201703cd702b4bb5e3cfc4efddaa4e2b4c58fce3703baa456a11e +size 1422130776 diff --git a/gemma4_gguf/temp_split_0de139d7_001.safetensors b/gemma4_gguf/temp_split_0de139d7_001.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fef27bf46d1d0340526be861b28df39b40ccd6ab --- /dev/null +++ b/gemma4_gguf/temp_split_0de139d7_001.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2ea02bad44856b48a53786446e9d32217ea9c377b26d16032919b8785c2b87be +size 4697620648 diff --git a/gemma4_gguf/temp_split_0de139d7_002.safetensors b/gemma4_gguf/temp_split_0de139d7_002.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ec2650695d5f1e7b37d435fc57fdba18cf1beb9c --- /dev/null +++ b/gemma4_gguf/temp_split_0de139d7_002.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85cebd4f098a5d61a725613814958869c645b022e6a7c083e69404732cd541b4 +size 1585731424 diff --git a/gemma4_gguf/temp_split_0de139d7_003.safetensors b/gemma4_gguf/temp_split_0de139d7_003.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..16ccab5ef52b5e0dd29520230c669843d904e881 --- /dev/null +++ b/gemma4_gguf/temp_split_0de139d7_003.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e169e6e4b22e2ab8787a697a0cb82391d5cd15dccdc183ee0053976a4bbf2375 +size 1446899712 diff --git a/gemma4_gguf/tokenizer.json b/gemma4_gguf/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..1ff9f3e3439a939b971f9919e821bf87e835a503 --- /dev/null +++ b/gemma4_gguf/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f +size 32169626 diff --git a/gemma4_gguf/tokenizer_config.json b/gemma4_gguf/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..91438f340be28997de476f607aff109b9f68f8b8 --- /dev/null +++ b/gemma4_gguf/tokenizer_config.json @@ -0,0 +1,290 @@ +{ + "audio_token": "<|audio|>", + "backend": "tokenizers", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "bos_token": "", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eos_token": "", + "eot_token": "", + "escape_token": "<|\"|>", + "etc_token": "", + "etd_token": "", + "etr_token": "", + "extra_special_tokens": [ + "<|video|>" + ], + "image_token": "<|image|>", + "is_local": false, + "mask_token": "", + "model_max_length": 131072, + "model_specific_special_tokens": { + "audio_token": "<|audio|>", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eot_token": "", + "escape_token": "<|\"|>", + "etc_token": "", + "etd_token": "", + "etr_token": "", + "image_token": "<|image|>", + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>" + }, + "pad_token": "", + "padding_side": "right", + "processor_class": "Gemma4Processor", + "response_schema": { + "properties": { + "content": { + "type": "string" + }, + "role": { + "const": "assistant" + }, + "thinking": { + "type": "string" + }, + "tool_calls": { + "items": { + "properties": { + "function": { + "properties": { + "arguments": { + "additionalProperties": {}, + "type": "object", + "x-parser": "gemma4-tool-call" + }, + "name": { + "type": "string" + } + }, + "type": "object", + "x-regex": "call\\:(?P\\w+)(?P\\{.*\\})" + }, + "type": { + "const": "function" + } + }, + "type": "object" + }, + "type": "array", + "x-regex-iterator": "<\\|tool_call>(.*?)" + } + }, + "type": "object", + "x-regex": "(\\<\\|channel\\>thought\\n(?P.*?)\\)?(?P\\<\\|tool_call\\>.*\\)?(?P(?:(?!\\)(?!\\<\\|tool_response\\>).)+)?(?:\\|\\<\\|tool_response\\>)?" + }, + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>", + "tokenizer_class": "GemmaTokenizer", + "unk_token": "", + "added_tokens_decoder": { + "0": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "1": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "2": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "3": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "4": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "46": { + "content": "<|tool>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "47": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "48": { + "content": "<|tool_call>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "49": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "50": { + "content": "<|tool_response>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "51": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "52": { + "content": "<|\"|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "98": { + "content": "<|think|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "100": { + "content": "<|channel>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "101": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "105": { + "content": "<|turn>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "106": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "255999": { + "content": "<|image>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "256000": { + "content": "<|audio>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258880": { + "content": "<|image|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258881": { + "content": "<|audio|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258882": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258883": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258884": { + "content": "<|video|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + }, + "chat_template": "{{ bos_token }}{%- macro strip_thinking(text) -%}\n {%- set ns = namespace(result='') -%}\n {%- for part in text.split('') -%}\n {%- if '<|channel>' in part -%}\n {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}\n {%- else -%}\n {%- set ns.result = ns.result + part -%}\n {%- endif -%}\n {%- endfor -%}\n {{- ns.result | trim -}}\n{%- endmacro -%}\n{%- set thinking = enable_thinking is defined and enable_thinking -%}\n{%- set loop_messages = messages -%}\n{%- if messages[0]['role'] in ['system', 'developer'] or thinking -%}\n {{ '<|turn>system\n' }}\n {%- if thinking -%}\n {{ '<|think|>\n' }}\n {%- endif -%}\n {%- if messages[0]['role'] in ['system', 'developer'] -%}\n {{ messages[0]['content'] | trim }}\n {%- set loop_messages = messages[1:] -%}\n {%- endif -%}\n {{ '\n' }}\n{%- endif -%}\n{%- for message in loop_messages -%}\n {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}\n {{ raise_exception(\"Conversation roles must alternate user/assistant/user/assistant/...\") }}\n {%- endif -%}\n {%- if (message['role'] == 'assistant') -%}\n {%- set role = \"model\" -%}\n {%- else -%}\n {%- set role = message['role'] -%}\n {%- endif -%}\n {{ '<|turn>' + role + '\n' }}\n {%- if message['content'] is string -%}\n {%- if role == \"model\" -%}\n {{ strip_thinking(message['content']) }}\n {%- else -%}\n {{ message['content'] | trim }}\n {%- endif -%}\n {%- elif message['content'] is iterable -%}\n {%- for item in message['content'] -%}\n {%- if item['type'] == 'audio' -%}\n {{ '<|audio|>' }}\n {%- elif item['type'] == 'image' -%}\n {{ '<|image|>' }}\n {%- elif item['type'] == 'video' -%}\n {{ '<|video|>' }}\n {%- elif item['type'] == 'text' -%}\n {%- if role == \"model\" -%}\n {{ strip_thinking(item['text']) }}\n {%- else -%}\n {{ item['text'] | trim }}\n {%- endif -%}\n {%- endif -%}\n {%- endfor -%}\n {%- else -%}\n {{ raise_exception(\"Invalid content type\") }}\n {%- endif -%}\n {{ '\n' }}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n {{'<|turn>model\n'}}\n{%- endif -%}\n" +} \ No newline at end of file diff --git a/lora_weights_final/README.md b/lora_weights_final/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7a2628dee0e8593d629689f1d074cefe86c35da5 --- /dev/null +++ b/lora_weights_final/README.md @@ -0,0 +1,210 @@ +--- +base_model: unsloth/gemma-4-e2b-it-unsloth-bnb-4bit +library_name: peft +pipeline_tag: text-generation +tags: +- base_model:adapter:unsloth/gemma-4-e2b-it-unsloth-bnb-4bit +- lora +- sft +- transformers +- trl +- unsloth +--- + +# Model Card for Model ID + + + + + +## Model Details + +### Model Description + + + + + +- **Developed by:** [More Information Needed] +- **Funded by [optional]:** [More Information Needed] +- **Shared by [optional]:** [More Information Needed] +- **Model type:** [More Information Needed] +- **Language(s) (NLP):** [More Information Needed] +- **License:** [More Information Needed] +- **Finetuned from model [optional]:** [More Information Needed] + +### Model Sources [optional] + + + +- **Repository:** [More Information Needed] +- **Paper [optional]:** [More Information Needed] +- **Demo [optional]:** [More Information Needed] + +## Uses + + + +### Direct Use + + + +[More Information Needed] + +### Downstream Use [optional] + + + +[More Information Needed] + +### Out-of-Scope Use + + + +[More Information Needed] + +## Bias, Risks, and Limitations + + + +[More Information Needed] + +### Recommendations + + + +Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations. + +## How to Get Started with the Model + +Use the code below to get started with the model. + +[More Information Needed] + +## Training Details + +### Training Data + + + +[More Information Needed] + +### Training Procedure + + + +#### Preprocessing [optional] + +[More Information Needed] + + +#### Training Hyperparameters + +- **Training regime:** [More Information Needed] + +#### Speeds, Sizes, Times [optional] + + + +[More Information Needed] + +## Evaluation + + + +### Testing Data, Factors & Metrics + +#### Testing Data + + + +[More Information Needed] + +#### Factors + + + +[More Information Needed] + +#### Metrics + + + +[More Information Needed] + +### Results + +[More Information Needed] + +#### Summary + + + +## Model Examination [optional] + + + +[More Information Needed] + +## Environmental Impact + + + +Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700). + +- **Hardware Type:** [More Information Needed] +- **Hours used:** [More Information Needed] +- **Cloud Provider:** [More Information Needed] +- **Compute Region:** [More Information Needed] +- **Carbon Emitted:** [More Information Needed] + +## Technical Specifications [optional] + +### Model Architecture and Objective + +[More Information Needed] + +### Compute Infrastructure + +[More Information Needed] + +#### Hardware + +[More Information Needed] + +#### Software + +[More Information Needed] + +## Citation [optional] + + + +**BibTeX:** + +[More Information Needed] + +**APA:** + +[More Information Needed] + +## Glossary [optional] + + + +[More Information Needed] + +## More Information [optional] + +[More Information Needed] + +## Model Card Authors [optional] + +[More Information Needed] + +## Model Card Contact + +[More Information Needed] +### Framework versions + +- PEFT 0.19.1 \ No newline at end of file diff --git a/lora_weights_final/adapter_config.json b/lora_weights_final/adapter_config.json new file mode 100644 index 0000000000000000000000000000000000000000..e1e34b7fd837cc424e8c00434a330db773803dca --- /dev/null +++ b/lora_weights_final/adapter_config.json @@ -0,0 +1,44 @@ +{ + "alora_invocation_tokens": null, + "alpha_pattern": {}, + "arrow_config": null, + "auto_mapping": { + "base_model_class": "Gemma4ForConditionalGeneration", + "parent_library": "transformers.models.gemma4.modeling_gemma4", + "unsloth_fixed": true + }, + "base_model_name_or_path": "unsloth/gemma-4-e2b-it-unsloth-bnb-4bit", + "bias": "none", + "corda_config": null, + "ensure_weight_tying": false, + "eva_config": null, + "exclude_modules": null, + "fan_in_fan_out": false, + "inference_mode": true, + "init_lora_weights": true, + "layer_replication": null, + "layers_pattern": null, + "layers_to_transform": null, + "loftq_config": {}, + "lora_alpha": 32, + "lora_bias": false, + "lora_dropout": 0, + "lora_ga_config": null, + "megatron_config": null, + "megatron_core": "megatron.core", + "modules_to_save": null, + "peft_type": "LORA", + "peft_version": "0.19.1", + "qalora_group_size": 16, + "r": 32, + "rank_pattern": {}, + "revision": null, + "target_modules": "(?:.*?(?:language|text).*?(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer).*?(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))|(?:\\bmodel\\.layers\\.[\\d]{1,}\\.(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer)\\.(?:(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)))", + "target_parameters": null, + "task_type": "CAUSAL_LM", + "trainable_token_indices": null, + "use_bdlora": null, + "use_dora": false, + "use_qalora": false, + "use_rslora": false +} \ No newline at end of file diff --git a/lora_weights_final/adapter_model.safetensors b/lora_weights_final/adapter_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a5b8be875ba039719502c80669f4e4c4267c7e94 --- /dev/null +++ b/lora_weights_final/adapter_model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f14e7d5c4c6cfe03e5f99ffea455059061c7edc6b443e909b0c236b0bc8913cb +size 101424416 diff --git a/lora_weights_final/chat_template.jinja b/lora_weights_final/chat_template.jinja new file mode 100644 index 0000000000000000000000000000000000000000..582bb4df95a62b60d248387fac933343a5878f83 --- /dev/null +++ b/lora_weights_final/chat_template.jinja @@ -0,0 +1,70 @@ +{{ bos_token }}{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} +{%- set thinking = enable_thinking is defined and enable_thinking -%} +{%- set loop_messages = messages -%} +{%- if messages[0]['role'] in ['system', 'developer'] or thinking -%} + {{ '<|turn>system +' }} + {%- if thinking -%} + {{ '<|think|> +' }} + {%- endif -%} + {%- if messages[0]['role'] in ['system', 'developer'] -%} + {{ messages[0]['content'] | trim }} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {{ ' +' }} +{%- endif -%} +{%- for message in loop_messages -%} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%} + {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }} + {%- endif -%} + {%- if (message['role'] == 'assistant') -%} + {%- set role = "model" -%} + {%- else -%} + {%- set role = message['role'] -%} + {%- endif -%} + {{ '<|turn>' + role + ' +' }} + {%- if message['content'] is string -%} + {%- if role == "model" -%} + {{ strip_thinking(message['content']) }} + {%- else -%} + {{ message['content'] | trim }} + {%- endif -%} + {%- elif message['content'] is iterable -%} + {%- for item in message['content'] -%} + {%- if item['type'] == 'audio' -%} + {{ '<|audio|>' }} + {%- elif item['type'] == 'image' -%} + {{ '<|image|>' }} + {%- elif item['type'] == 'video' -%} + {{ '<|video|>' }} + {%- elif item['type'] == 'text' -%} + {%- if role == "model" -%} + {{ strip_thinking(item['text']) }} + {%- else -%} + {{ item['text'] | trim }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{ raise_exception("Invalid content type") }} + {%- endif -%} + {{ ' +' }} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{'<|turn>model +'}} +{%- endif -%} diff --git a/lora_weights_final/tokenizer.json b/lora_weights_final/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..1ff9f3e3439a939b971f9919e821bf87e835a503 --- /dev/null +++ b/lora_weights_final/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f +size 32169626 diff --git a/lora_weights_final/tokenizer_config.json b/lora_weights_final/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..24b1522bb32b0eb05b631b1a9f112751b8cb0fab --- /dev/null +++ b/lora_weights_final/tokenizer_config.json @@ -0,0 +1,289 @@ +{ + "audio_token": "<|audio|>", + "backend": "tokenizers", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "bos_token": "", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eos_token": "", + "eot_token": "", + "escape_token": "<|\"|>", + "etc_token": "", + "etd_token": "", + "etr_token": "", + "extra_special_tokens": [ + "<|video|>" + ], + "image_token": "<|image|>", + "is_local": false, + "mask_token": "", + "model_max_length": 131072, + "model_specific_special_tokens": { + "audio_token": "<|audio|>", + "boa_token": "<|audio>", + "boi_token": "<|image>", + "eoa_token": "", + "eoc_token": "", + "eoi_token": "", + "eot_token": "", + "escape_token": "<|\"|>", + "etc_token": "", + "etd_token": "", + "etr_token": "", + "image_token": "<|image|>", + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>" + }, + "pad_token": "", + "padding_side": "right", + "processor_class": "Gemma4Processor", + "response_schema": { + "properties": { + "content": { + "type": "string" + }, + "role": { + "const": "assistant" + }, + "thinking": { + "type": "string" + }, + "tool_calls": { + "items": { + "properties": { + "function": { + "properties": { + "arguments": { + "additionalProperties": {}, + "type": "object", + "x-parser": "gemma4-tool-call" + }, + "name": { + "type": "string" + } + }, + "type": "object", + "x-regex": "call\\:(?P\\w+)(?P\\{.*\\})" + }, + "type": { + "const": "function" + } + }, + "type": "object" + }, + "type": "array", + "x-regex-iterator": "<\\|tool_call>(.*?)" + } + }, + "type": "object", + "x-regex": "(\\<\\|channel\\>thought\\n(?P.*?)\\)?(?P\\<\\|tool_call\\>.*\\)?(?P(?:(?!\\)(?!\\<\\|tool_response\\>).)+)?(?:\\|\\<\\|tool_response\\>)?" + }, + "soc_token": "<|channel>", + "sot_token": "<|turn>", + "stc_token": "<|tool_call>", + "std_token": "<|tool>", + "str_token": "<|tool_response>", + "think_token": "<|think|>", + "tokenizer_class": "GemmaTokenizer", + "unk_token": "", + "added_tokens_decoder": { + "0": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "1": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "2": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "3": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "4": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "46": { + "content": "<|tool>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "47": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "48": { + "content": "<|tool_call>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "49": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "50": { + "content": "<|tool_response>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "51": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "52": { + "content": "<|\"|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "98": { + "content": "<|think|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "100": { + "content": "<|channel>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "101": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "105": { + "content": "<|turn>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "106": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "255999": { + "content": "<|image>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "256000": { + "content": "<|audio>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258880": { + "content": "<|image|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258881": { + "content": "<|audio|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258882": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258883": { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + "258884": { + "content": "<|video|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + } +} diff --git a/unsloth_compiled_cache/AqlmLoraLinear_peft_forward.py b/unsloth_compiled_cache/AqlmLoraLinear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..f1db185980b5c7b4dc6b4118680b587e9b96edeb --- /dev/null +++ b/unsloth_compiled_cache/AqlmLoraLinear_peft_forward.py @@ -0,0 +1,89 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.aqlm import (torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor): + # note: logic differs from default Linear because merging is not supported + result = self.base_layer(x) + + if self.disable_adapters: + return result + + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + x = self._cast_input_dtype(x, lora_A.weight.dtype) + + output = lora_B(lora_A(dropout(x))) + if requires_conversion: + output = output.to(expected_dtype) + output = output * scaling + result += output + return result diff --git a/unsloth_compiled_cache/AwqLoraLinear_peft_forward.py b/unsloth_compiled_cache/AwqLoraLinear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..63183f741b34915002268815bac8aa3e0143f49c --- /dev/null +++ b/unsloth_compiled_cache/AwqLoraLinear_peft_forward.py @@ -0,0 +1,88 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.awq import (torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor): + result = self.quant_linear_module(x) + + if self.disable_adapters: + return result + + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + x = self._cast_input_dtype(x, lora_A.weight.dtype) + + output = lora_B(lora_A(dropout(x))) + if requires_conversion: + output = output.to(expected_dtype) + output = output * scaling + result = result + output + return result diff --git a/unsloth_compiled_cache/BatchNorm1d.py b/unsloth_compiled_cache/BatchNorm1d.py new file mode 100644 index 0000000000000000000000000000000000000000..bb363c9627ebb52c6c841a7e1d42815e86e948cf --- /dev/null +++ b/unsloth_compiled_cache/BatchNorm1d.py @@ -0,0 +1,121 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, nn) + +def forward(self, input: Tensor) -> Tensor: + self._check_input_dim(input) + + # exponential_average_factor is set to self.momentum + # (when it is available) only so that it gets updated + # in ONNX graph when this node is exported to ONNX. + if self.momentum is None: + exponential_average_factor = 0.0 + else: + exponential_average_factor = self.momentum + + if self.training and self.track_running_stats: + # TODO: if statement only here to tell the jit to skip emitting this when it is None + if self.num_batches_tracked is not None: # type: ignore[has-type] + self.num_batches_tracked.add_(1) # type: ignore[has-type] + if self.momentum is None: # use cumulative moving average + exponential_average_factor = 1.0 / float(self.num_batches_tracked) + else: # use exponential moving average + exponential_average_factor = self.momentum + + r""" + Decide whether the mini-batch stats should be used for normalization rather than the buffers. + Mini-batch stats are used in training mode, and in eval mode when buffers are None. + """ + if self.training: + bn_training = True + else: + bn_training = (self.running_mean is None) and (self.running_var is None) + + r""" + Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be + passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are + used for normalization (i.e. in eval mode when buffers are not None). + """ + return F.batch_norm( + input, + # If buffers are not to be tracked, ensure that they won't be updated + ( + self.running_mean + if not self.training or self.track_running_stats + else None + ), + self.running_var if not self.training or self.track_running_stats else None, + self.weight, + self.bias, + bn_training, + exponential_average_factor, + self.eps, + ).to(input.dtype).to(input.dtype) diff --git a/unsloth_compiled_cache/BatchNorm2d.py b/unsloth_compiled_cache/BatchNorm2d.py new file mode 100644 index 0000000000000000000000000000000000000000..bb363c9627ebb52c6c841a7e1d42815e86e948cf --- /dev/null +++ b/unsloth_compiled_cache/BatchNorm2d.py @@ -0,0 +1,121 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, nn) + +def forward(self, input: Tensor) -> Tensor: + self._check_input_dim(input) + + # exponential_average_factor is set to self.momentum + # (when it is available) only so that it gets updated + # in ONNX graph when this node is exported to ONNX. + if self.momentum is None: + exponential_average_factor = 0.0 + else: + exponential_average_factor = self.momentum + + if self.training and self.track_running_stats: + # TODO: if statement only here to tell the jit to skip emitting this when it is None + if self.num_batches_tracked is not None: # type: ignore[has-type] + self.num_batches_tracked.add_(1) # type: ignore[has-type] + if self.momentum is None: # use cumulative moving average + exponential_average_factor = 1.0 / float(self.num_batches_tracked) + else: # use exponential moving average + exponential_average_factor = self.momentum + + r""" + Decide whether the mini-batch stats should be used for normalization rather than the buffers. + Mini-batch stats are used in training mode, and in eval mode when buffers are None. + """ + if self.training: + bn_training = True + else: + bn_training = (self.running_mean is None) and (self.running_var is None) + + r""" + Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be + passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are + used for normalization (i.e. in eval mode when buffers are not None). + """ + return F.batch_norm( + input, + # If buffers are not to be tracked, ensure that they won't be updated + ( + self.running_mean + if not self.training or self.track_running_stats + else None + ), + self.running_var if not self.training or self.track_running_stats else None, + self.weight, + self.bias, + bn_training, + exponential_average_factor, + self.eps, + ).to(input.dtype).to(input.dtype) diff --git a/unsloth_compiled_cache/BatchNorm3d.py b/unsloth_compiled_cache/BatchNorm3d.py new file mode 100644 index 0000000000000000000000000000000000000000..bb363c9627ebb52c6c841a7e1d42815e86e948cf --- /dev/null +++ b/unsloth_compiled_cache/BatchNorm3d.py @@ -0,0 +1,121 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, nn) + +def forward(self, input: Tensor) -> Tensor: + self._check_input_dim(input) + + # exponential_average_factor is set to self.momentum + # (when it is available) only so that it gets updated + # in ONNX graph when this node is exported to ONNX. + if self.momentum is None: + exponential_average_factor = 0.0 + else: + exponential_average_factor = self.momentum + + if self.training and self.track_running_stats: + # TODO: if statement only here to tell the jit to skip emitting this when it is None + if self.num_batches_tracked is not None: # type: ignore[has-type] + self.num_batches_tracked.add_(1) # type: ignore[has-type] + if self.momentum is None: # use cumulative moving average + exponential_average_factor = 1.0 / float(self.num_batches_tracked) + else: # use exponential moving average + exponential_average_factor = self.momentum + + r""" + Decide whether the mini-batch stats should be used for normalization rather than the buffers. + Mini-batch stats are used in training mode, and in eval mode when buffers are None. + """ + if self.training: + bn_training = True + else: + bn_training = (self.running_mean is None) and (self.running_var is None) + + r""" + Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be + passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are + used for normalization (i.e. in eval mode when buffers are not None). + """ + return F.batch_norm( + input, + # If buffers are not to be tracked, ensure that they won't be updated + ( + self.running_mean + if not self.training or self.track_running_stats + else None + ), + self.running_var if not self.training or self.track_running_stats else None, + self.weight, + self.bias, + bn_training, + exponential_average_factor, + self.eps, + ).to(input.dtype).to(input.dtype) diff --git a/unsloth_compiled_cache/BlockDiagonalLinear_peft_forward.py b/unsloth_compiled_cache/BlockDiagonalLinear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..023888df2797eeec263ef4ccb0eedfeeb9a79dfb --- /dev/null +++ b/unsloth_compiled_cache/BlockDiagonalLinear_peft_forward.py @@ -0,0 +1,75 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.variants import (torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor) -> torch.Tensor: + first_dims = x.shape[:-1] + if x.dim() != 2: + x = x.reshape(-1, x.shape[-1]) + B = x.shape[0] + nb = self.nblocks + m = x.shape[-1] // nb + n = self.out_features // nb + x = x.reshape(B, nb, m) + w = self.weight.view(nb, n, m) + out = torch.einsum("bim,inm->bin", x, w) + return out.reshape(*first_dims, -1) diff --git a/unsloth_compiled_cache/Conv1d.py b/unsloth_compiled_cache/Conv1d.py new file mode 100644 index 0000000000000000000000000000000000000000..6bebb35f306c9cda88a06951799f3172b9f4fcae --- /dev/null +++ b/unsloth_compiled_cache/Conv1d.py @@ -0,0 +1,78 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable + + +def forward(self, input: Tensor) -> Tensor: + original_dtype = input.dtype + input = input.to(self.weight.dtype) + original_dtype = input.dtype + input = input.to(self.weight.dtype) + return self._conv_forward(input, self.weight, self.bias).to(original_dtype).to(original_dtype) diff --git a/unsloth_compiled_cache/Conv2d.py b/unsloth_compiled_cache/Conv2d.py new file mode 100644 index 0000000000000000000000000000000000000000..6bebb35f306c9cda88a06951799f3172b9f4fcae --- /dev/null +++ b/unsloth_compiled_cache/Conv2d.py @@ -0,0 +1,78 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable + + +def forward(self, input: Tensor) -> Tensor: + original_dtype = input.dtype + input = input.to(self.weight.dtype) + original_dtype = input.dtype + input = input.to(self.weight.dtype) + return self._conv_forward(input, self.weight, self.bias).to(original_dtype).to(original_dtype) diff --git a/unsloth_compiled_cache/Conv3d.py b/unsloth_compiled_cache/Conv3d.py new file mode 100644 index 0000000000000000000000000000000000000000..6bebb35f306c9cda88a06951799f3172b9f4fcae --- /dev/null +++ b/unsloth_compiled_cache/Conv3d.py @@ -0,0 +1,78 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable + + +def forward(self, input: Tensor) -> Tensor: + original_dtype = input.dtype + input = input.to(self.weight.dtype) + original_dtype = input.dtype + input = input.to(self.weight.dtype) + return self._conv_forward(input, self.weight, self.bias).to(original_dtype).to(original_dtype) diff --git a/unsloth_compiled_cache/ConvTranspose1d.py b/unsloth_compiled_cache/ConvTranspose1d.py new file mode 100644 index 0000000000000000000000000000000000000000..6638c100b7e775ceb5a2b029fd378919b791f52e --- /dev/null +++ b/unsloth_compiled_cache/ConvTranspose1d.py @@ -0,0 +1,105 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, nn) + +def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor: + original_dtype = input.dtype + input = input.to(self.weight.dtype) + original_dtype = input.dtype + input = input.to(self.weight.dtype) + if self.padding_mode != "zeros": + raise ValueError( + "Only `zeros` padding mode is supported for ConvTranspose1d" + ) + + assert isinstance(self.padding, tuple) + # One cannot replace List by Tuple or Sequence in "_output_padding" because + # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`. + num_spatial_dims = 1 + output_padding = self._output_padding( + input, + output_size, + self.stride, # type: ignore[arg-type] + self.padding, # type: ignore[arg-type] + self.kernel_size, # type: ignore[arg-type] + num_spatial_dims, + self.dilation, # type: ignore[arg-type] + ) + return F.conv_transpose1d( + input, + self.weight, + self.bias, + self.stride, + self.padding, + output_padding, + self.groups, + self.dilation, + ).to(original_dtype).to(original_dtype) diff --git a/unsloth_compiled_cache/ConvTranspose2d.py b/unsloth_compiled_cache/ConvTranspose2d.py new file mode 100644 index 0000000000000000000000000000000000000000..1a7a3faffcb0fbcb499d87a29183c68efdab7801 --- /dev/null +++ b/unsloth_compiled_cache/ConvTranspose2d.py @@ -0,0 +1,114 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, nn) + +def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor: + original_dtype = input.dtype + input = input.to(self.weight.dtype) + original_dtype = input.dtype + input = input.to(self.weight.dtype) + """ + Performs the forward pass. + + Attributes: + input (Tensor): The input tensor. + output_size (list[int], optional): A list of integers representing + the size of the output tensor. Default is None. + """ + if self.padding_mode != "zeros": + raise ValueError( + "Only `zeros` padding mode is supported for ConvTranspose2d" + ) + + assert isinstance(self.padding, tuple) + # One cannot replace List by Tuple or Sequence in "_output_padding" because + # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`. + num_spatial_dims = 2 + output_padding = self._output_padding( + input, + output_size, + self.stride, # type: ignore[arg-type] + self.padding, # type: ignore[arg-type] + self.kernel_size, # type: ignore[arg-type] + num_spatial_dims, + self.dilation, # type: ignore[arg-type] + ) + + return F.conv_transpose2d( + input, + self.weight, + self.bias, + self.stride, + self.padding, + output_padding, + self.groups, + self.dilation, + ).to(original_dtype).to(original_dtype) diff --git a/unsloth_compiled_cache/ConvTranspose3d.py b/unsloth_compiled_cache/ConvTranspose3d.py new file mode 100644 index 0000000000000000000000000000000000000000..db25cd2316c03c1af179d7dd6a39a2cf7d81210b --- /dev/null +++ b/unsloth_compiled_cache/ConvTranspose3d.py @@ -0,0 +1,106 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, nn) + +def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor: + original_dtype = input.dtype + input = input.to(self.weight.dtype) + original_dtype = input.dtype + input = input.to(self.weight.dtype) + if self.padding_mode != "zeros": + raise ValueError( + "Only `zeros` padding mode is supported for ConvTranspose3d" + ) + + assert isinstance(self.padding, tuple) + # One cannot replace List by Tuple or Sequence in "_output_padding" because + # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`. + num_spatial_dims = 3 + output_padding = self._output_padding( + input, + output_size, + self.stride, # type: ignore[arg-type] + self.padding, # type: ignore[arg-type] + self.kernel_size, # type: ignore[arg-type] + num_spatial_dims, + self.dilation, # type: ignore[arg-type] + ) + + return F.conv_transpose3d( + input, + self.weight, + self.bias, + self.stride, + self.padding, + output_padding, + self.groups, + self.dilation, + ).to(original_dtype).to(original_dtype) diff --git a/unsloth_compiled_cache/GPTQLoraLinear_peft_forward.py b/unsloth_compiled_cache/GPTQLoraLinear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..25d190b5fb7431376b82932a885e323ecacfe7bc --- /dev/null +++ b/unsloth_compiled_cache/GPTQLoraLinear_peft_forward.py @@ -0,0 +1,96 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.gptq import (torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor): + # note: logic differs from default Linear because merging is not supported + result = self.quant_linear_module(x) + + if self.disable_adapters: + return result + + lora_A_keys = self.lora_A.keys() + + for active_adapter in self.active_adapters: + if active_adapter not in lora_A_keys: + continue + torch_result_dtype = result.dtype + + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + + if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) + + if active_adapter not in self.lora_variant: # vanilla LoRA + return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) + else: + result = self.lora_variant[active_adapter].forward( + self, + active_adapter=active_adapter, + x=x, + result=result, + ) + + result = result.to(torch_result_dtype) + return result diff --git a/unsloth_compiled_cache/GroupNorm.py b/unsloth_compiled_cache/GroupNorm.py new file mode 100644 index 0000000000000000000000000000000000000000..fddaa750508343174e4cfc5b6aa91c34dac5db4c --- /dev/null +++ b/unsloth_compiled_cache/GroupNorm.py @@ -0,0 +1,74 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F) + +def forward(self, input: Tensor) -> Tensor: + return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps).to(input.dtype).to(input.dtype) diff --git a/unsloth_compiled_cache/LayerNorm.py b/unsloth_compiled_cache/LayerNorm.py new file mode 100644 index 0000000000000000000000000000000000000000..964cd3e649ca443490ee29ffa442fd71e971d0ad --- /dev/null +++ b/unsloth_compiled_cache/LayerNorm.py @@ -0,0 +1,76 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F) + +def forward(self, input: Tensor) -> Tensor: + return F.layer_norm( + input, self.normalized_shape, self.weight, self.bias, self.eps + ).to(input.dtype).to(input.dtype) diff --git a/unsloth_compiled_cache/Linear4bit_peft_forward.py b/unsloth_compiled_cache/Linear4bit_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..96bd71e51defce447c52ed6128b031db57b39fec --- /dev/null +++ b/unsloth_compiled_cache/Linear4bit_peft_forward.py @@ -0,0 +1,126 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +try: + from peft.tuners.lora.layer import VARIANT_KWARG_KEYS +except ImportError: + VARIANT_KWARG_KEYS = ['alora_offsets'] +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.bnb import (VARIANT_KWARG_KEYS, torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + + adapter_names = kwargs.pop("adapter_names", None) + variant_kwargs = {k: kwargs.pop(k, None) for k in VARIANT_KWARG_KEYS} # don't pass these to base_layer + + if self.disable_adapters: + if self.merged: + self.unmerge() + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + elif adapter_names is not None: + result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **variant_kwargs, **kwargs) + elif self.merged: + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + else: + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + # As per Tim Dettmers, for 4bit, we need to defensively clone here. + # The reason is that in some cases, an error can occur that backprop + # does not work on a manipulated view. This issue may be solved with + # newer PyTorch versions but this would need extensive testing to be + # sure. + + + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + x = self._cast_input_dtype(x, lora_A.weight.dtype) + + if active_adapter not in self.lora_variant: # vanilla LoRA + return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(result.dtype) + if requires_conversion: + output = output.to(expected_dtype) + result = result + output + else: + result = self.lora_variant[active_adapter].forward( + self, + active_adapter=active_adapter, + x=x, + result=result, + **variant_kwargs, + **kwargs, + ) + if requires_conversion: + result = result.to(expected_dtype) + + return result diff --git a/unsloth_compiled_cache/Linear8bitLt_peft_forward.py b/unsloth_compiled_cache/Linear8bitLt_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..fb1a296bce8af2c3969d23fc9df21e50e5b99b83 --- /dev/null +++ b/unsloth_compiled_cache/Linear8bitLt_peft_forward.py @@ -0,0 +1,118 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +try: + from peft.tuners.lora.layer import VARIANT_KWARG_KEYS +except ImportError: + VARIANT_KWARG_KEYS = ['alora_offsets'] +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} + +import torch._dynamo +@torch._dynamo.disable +def _call_8bit_base_layer(base_layer, x, *args, **kwargs): + return base_layer(x, *args, **kwargs) +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.bnb import (VARIANT_KWARG_KEYS, torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + + adapter_names = kwargs.pop("adapter_names", None) + variant_kwargs = {k: kwargs.pop(k, None) for k in VARIANT_KWARG_KEYS} # don't pass these to base_layer + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = _call_8bit_base_layer(self.base_layer, x, *args, **kwargs) + elif adapter_names is not None: + result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **variant_kwargs, **kwargs) + elif self.merged: + result = _call_8bit_base_layer(self.base_layer, x, *args, **kwargs) + else: + result = _call_8bit_base_layer(self.base_layer, x, *args, **kwargs) + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + x = self._cast_input_dtype(x, lora_A.weight.dtype) + + if active_adapter not in self.lora_variant: # vanilla LoRA + return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(result.dtype) + if requires_conversion: + output = output.to(expected_dtype) + result = result + output + else: + result = self.lora_variant[active_adapter].forward( + self, + active_adapter=active_adapter, + x=x, + result=result, + **variant_kwargs, + **kwargs, + ) + if requires_conversion: + result = result.to(expected_dtype) + + return result diff --git a/unsloth_compiled_cache/Linear_peft_forward.py b/unsloth_compiled_cache/Linear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..f8ea61192c840e0fff59309541784b3436e69cb8 --- /dev/null +++ b/unsloth_compiled_cache/Linear_peft_forward.py @@ -0,0 +1,115 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +try: + from peft.tuners.lora.layer import VARIANT_KWARG_KEYS +except ImportError: + VARIANT_KWARG_KEYS = ['alora_offsets'] +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.variants import (Any, torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + + adapter_names = kwargs.pop("adapter_names", None) + variant_kwargs = {k: kwargs.pop(k, None) for k in VARIANT_KWARG_KEYS} # don't pass these to base_layer + + if self.disable_adapters: + if self.merged: + self.unmerge() + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + elif adapter_names is not None: + result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **variant_kwargs, **kwargs) + elif self.merged: + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + else: + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + torch_result_dtype = result.dtype + + lora_A_keys = self.lora_A.keys() + for active_adapter in self.active_adapters: + if active_adapter not in lora_A_keys: + continue + + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) + if active_adapter not in self.lora_variant: # vanilla LoRA + return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) + else: + result = self.lora_variant[active_adapter].forward( + self, + active_adapter=active_adapter, + x=x, + result=result, + **variant_kwargs, + **kwargs, + ) + + result = result.to(torch_result_dtype) + + return result diff --git a/unsloth_compiled_cache/LoraParallelLinear_peft_forward.py b/unsloth_compiled_cache/LoraParallelLinear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..1c54d5776e5492d35c80d7f94386a1f91dcf5e09 --- /dev/null +++ b/unsloth_compiled_cache/LoraParallelLinear_peft_forward.py @@ -0,0 +1,92 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.tp_layer import (Any, __name__, torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): + + adapter_names = kwargs.pop("adapter_names", None) + # If weight is used for matrix multiplication here, the final aggregation operation of the original + # parallel_linear layer will be missing, so we need to directly call its forward function to obtain the + # output of the original parallel_linear layer. + if self.disable_adapters: + if self.merged: + self.unmerge() + result, bias = self.base_layer(x, *args, **kwargs) + elif adapter_names is not None: + raise ValueError(f"{self.__class__.__name__} does not support mixed_batch_forward yet.") + elif self.merged: + result, bias = self.base_layer(x, *args, **kwargs) + else: + result, bias = self.base_layer(x, *args, **kwargs) + torch_result_dtype = result.dtype + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) + return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) + + result = result.to(torch_result_dtype) + return result, bias diff --git a/unsloth_compiled_cache/RMSNorm.py b/unsloth_compiled_cache/RMSNorm.py new file mode 100644 index 0000000000000000000000000000000000000000..f2f4ab051012d8118e4ac840c29d1a32f08dc54c --- /dev/null +++ b/unsloth_compiled_cache/RMSNorm.py @@ -0,0 +1,77 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, torch) + +def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Runs the forward pass. + """ + return F.rms_norm(x, self.normalized_shape, self.weight, self.eps).to(x.dtype).to(x.dtype) diff --git a/unsloth_compiled_cache/TeLinear_peft_forward.py b/unsloth_compiled_cache/TeLinear_peft_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..fb1c017b4a5251646ff4338699f35b523dabd2b2 --- /dev/null +++ b/unsloth_compiled_cache/TeLinear_peft_forward.py @@ -0,0 +1,94 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from peft.tuners.lora.te import (Any, __name__, torch) + + +torch_addmm = torch.addmm +torch_add = torch.add +# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def lora_forward(result, lora_A, lora_B, dropout, x, scaling): + # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 + # by _cast_input_dtype when autocast is disabled + target_dtype = result.dtype + xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() + # output = result + scaling * xA @ lora_B.weight.t() + shape = result.shape + output = torch_addmm( + result.view(-1, shape[-1]), + xA.view(-1, xA.shape[-1]), + lora_B.weight.to(target_dtype).t(), + alpha = scaling, + beta = 1, + ).view(shape) + + bias = lora_B.bias + if bias is not None: + output = torch_add( + output, + bias.to(target_dtype), + alpha = scaling, + ) + return output +pass + +def unsloth_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + + adapter_names = kwargs.pop("adapter_names", None) + + if self.disable_adapters: + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + elif adapter_names is not None: + raise ValueError(f"{self.__class__.__name__} does not support mixed_batch_forward yet.") + else: + if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and x.dtype != self.base_layer.weight.dtype: + x = x.to(self.base_layer.weight.dtype) + result = self.base_layer(x, *args, **kwargs) + torch_result_dtype = result.dtype + + lora_A_keys = self.lora_A.keys() + for active_adapter in self.active_adapters: + if active_adapter not in lora_A_keys: + continue + + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) + return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) + + result = result.to(torch_result_dtype) + + return result diff --git a/unsloth_compiled_cache/UnslothDPOTrainer.py b/unsloth_compiled_cache/UnslothDPOTrainer.py new file mode 100644 index 0000000000000000000000000000000000000000..c1035b7eeb7b477c4a3923f37b2000ba96b307a5 --- /dev/null +++ b/unsloth_compiled_cache/UnslothDPOTrainer.py @@ -0,0 +1,2348 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.dpo_trainer import (Any, AutoProcessor, Callable, DPOConfig, DPOTrainer, DataCollator, DataCollatorForPreference, DataCollatorForVisionPreference, DataLoader, Dataset, EvalPrediction, F, Hasher, IterableDataset, IterableDatasetDict, LoraConfig, PartialState, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SyncRefModelCallback, TrainerCallback, Version, _BaseTrainer, apply_chat_template, concatenate_datasets, contextlib, create_model_from_path, dataclass, defaultdict, disable_dropout_in_model, disable_gradient_checkpointing, entropy_from_logits, extract_prompt, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, hash_module, is_conversational, is_liger_kernel_available, is_peft_available, is_peft_model, json, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, selective_log_softmax, textwrap, torch, tqdm, transformers, use_adapter, AutoProcessor, Callable, DPOConfig, DPOTrainer, DataCollator, DataCollatorForPreference, DataCollatorForVisionPreference, Dataset, EvalPrediction, F, IterableDataset, IterableDatasetDict, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SyncRefModelCallback, TrainerCallback, Version, contextlib, create_model_from_path, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, is_liger_kernel_available, is_peft_available, is_peft_model, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, torch, transformers, F, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch) + + +import os +import math +import logging +from typing import * +from dataclasses import dataclass, field +from packaging.version import Version +import torch +import numpy as np +from contextlib import nullcontext +from torch.nn import functional as F +import inspect +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling +from transformers.training_args import ParallelMode +from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize + +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass +# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain +# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. +try: + from unsloth.models._utils import _unsloth_reset_stray_compile_cache +except Exception: + def _unsloth_reset_stray_compile_cache(self): pass +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Drop any torch.compile graph cache poisoned by a stray pre-train forward. + try: + _unsloth_reset_stray_compile_cache(self) + except Exception: + pass + # Finish the previous W&B run if this is a subsequent train() call. + # We do this at the START of train() (not the end) so that + # evaluate() / log() still work after train() completes. + # HF's WandbCallback.setup() will call wandb.init() for the new run. + # See: https://github.com/unslothai/unsloth/issues/3954 + if getattr(self, '_unsloth_training_completed', False): + try: + import wandb + if wandb.run is not None: + wandb.finish() + # Reset HF's WandbCallback so it calls wandb.init() for the new run + for cb in self.callback_handler.callbacks: + if type(cb).__name__ == 'WandbCallback': + cb._initialized = False + break + except: + pass + # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + output = f(self, *args, **kwargs) + # Restore previous mode when possible + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass + # Mark that training completed so the next train() call can + # finish this W&B run before starting a new one + self._unsloth_training_completed = True + return output + return wrapper +pass + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_hidden_states_selective_log_softmax( + hidden_states: torch.Tensor, + lm_head: torch.Tensor, + index: torch.Tensor, + chunks: int = 4, + logit_scale_multiply: float = 0.0, + logit_scale_divide: float = 0.0, + logit_softcapping: float = 0.0, + temperature: float = 1.0, +) -> torch.Tensor: + # All Unsloth Zoo code licensed under AGPL3 + flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + flat_index = index.reshape(-1) + + chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) + chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) + + all_per_token_logps = [] + + for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): + chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() + + if logit_scale_multiply != 0.0: + chunk_logits = chunk_logits * logit_scale_multiply + if logit_scale_divide != 0.0: + chunk_logits = chunk_logits / logit_scale_divide + if logit_softcapping != 0.0: + chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) + + chunk_logits = chunk_logits.to(torch.float32) + + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + + selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + + all_per_token_logps = torch.concat(all_per_token_logps) + + all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) + return all_per_token_logps + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_selective_log_softmax( + logits, + index, + temperature: float = 1.0, + chunks: int = 4, +): + chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) + chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) + all_per_token_logps = [] + # Per-chunk selective_log_softmax. + for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): + chunk_logits = chunk_logits.to(torch.float32) + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + pass + all_per_token_logps = torch.concat(all_per_token_logps) + all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) + return all_per_token_logps + +def calculate_pad_tokens_in_prompt( + input_ids: torch.Tensor, + logits_to_keep: int, + pad_token_id: int +) -> torch.Tensor: + """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" + if logits_to_keep >= input_ids.shape[1]: + raise ValueError("logits_to_keep must be smaller than the sequence length.") + + prompt_section = input_ids[:, :-logits_to_keep] + + padding_mask = (prompt_section == pad_token_id) + + pad_token_counts = padding_mask.sum(dim=1) + + return pad_token_counts + +def create_completion_attention_mask( + completion_input_ids: torch.Tensor, + left_pad_tokens_per_prompt: torch.Tensor, + max_left_pad: int, + pad_token_id: int +) -> torch.Tensor: + """Build a completion mask that zeros leading prompt and trailing pad tokens. + + For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) + this returns [0,0,0,1,1,1,0,0,0]. + """ + batch_size, completion_len = completion_input_ids.shape + device = completion_input_ids.device + + num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt + + indices = torch.arange(completion_len, device=device).unsqueeze(0) + shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) + + non_padding_mask = (completion_input_ids != pad_token_id) + + final_mask = shift_mask & non_padding_mask + + return final_mask + +def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: + """Move all padding tokens in each sequence to the right.""" + mask = (tensor != pad_id) + # stable=True since the binary mask is unordered. + sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) + packed_tensor = torch.gather(tensor, 1, sorted_indices) + return packed_tensor + +def align_logprobs_with_mask( + logprob_tensor: torch.Tensor, + attention_mask: torch.Tensor, + pad_value: float = 0.0 +) -> torch.Tensor: + """Align a log probability tensor with a given attention mask.""" + + device = logprob_tensor.device + batch_size, logprob_seq_len = logprob_tensor.shape + mask_seq_len = attention_mask.shape[1] + + padded_logprobs = torch.full( + attention_mask.shape, + fill_value=pad_value, + dtype=logprob_tensor.dtype, + device=device + ) + + left_pad_counts = torch.argmax(attention_mask, dim=1) + + cols = torch.arange(logprob_seq_len, device=device) + dest_indices = left_pad_counts.unsqueeze(1) + cols + + # Destination row indices, shape [batch_size, logprob_seq_len]. + row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) + + # Keep only in-bounds destinations, then scatter via advanced indexing. + valid_mask = dest_indices < mask_seq_len + valid_rows = row_indices[valid_mask] + valid_cols = dest_indices[valid_mask] + valid_vals = logprob_tensor[valid_mask] + padded_logprobs[valid_rows, valid_cols] = valid_vals + + return padded_logprobs + +def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" + if tool_mask is None: + return completion_mask + if tool_mask.shape[0] != completion_mask.shape[0]: + raise ValueError("tool_mask batch size must match completion_mask batch size.") + + tool_mask = tool_mask.to(device=completion_mask.device) + if tool_mask.shape == completion_mask.shape: + aligned_tool_mask = tool_mask + else: + aligned_tool_mask = align_logprobs_with_mask( + tool_mask, + completion_mask, + pad_value=0, + ) + return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) + +def autotune_batch_and_chunks( + total_input_rows, + seq_len, + hidden_size, + vocab_size, + dtype_bytes=16, + multiplier=None +): + if multiplier is None: + final_m = max(4, seq_len // 4096) + else: + final_m = multiplier + + if torch.cuda.is_available(): + free_bytes, _ = torch.cuda.mem_get_info() + limit_gb = (free_bytes / (1024**3))*.80 + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + # XPU: estimate free memory as total - reserved. + total_mem = torch.xpu.get_device_properties(0).total_memory + reserved_mem = torch.xpu.memory_reserved() + free_bytes = total_mem - reserved_mem + limit_gb = (free_bytes / (1024**3)) * 0.80 + else: + # Fallback: assume 8GB available. + limit_gb = 8.0 + + bytes_to_gb = 1024**3 + + b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) + + hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb + + base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb + logits_gb = base_logits / final_m + + total_mem_gb = hidden_gb + logits_gb + + valid_mask = total_mem_gb <= limit_gb + valid_indices = torch.nonzero(valid_mask, as_tuple=False) + + if valid_indices.shape[0] == 0: + #This means your GPU will OOM + return 4, final_m + + best_idx = valid_indices[0].item() + final_b = int(b_vals[best_idx].item()) + + return final_b, final_m + +def sanitize_logprob(logprob): + """Local port of trl.scripts.vllm_serve.sanitize_logprob. + Filters NaN logprobs from vLLM outputs.""" + value = logprob.logprob + if math.isnan(value): + logging.getLogger(__name__).warning( + f"Generated NaN logprob, token logprob '{logprob}' will be ignored" + ) + return None + return value +def dpo_trainer_vision_process_row( + features, + processing_class, + max_prompt_length = None, + max_completion_length = None, + add_special_tokens = True, + is_chat = False, +): + text = features.get("prompt", "") + images = features.get("images") + processor, tokenizer = processing_class, processing_class.tokenizer + processed_features = processor( + images = images, + text = text, + add_special_tokens = False, + ) + + prompt_input_ids = processed_features["input_ids"][0] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens = False)["input_ids"] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens = False)["input_ids"] + + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + if not is_chat and tokenizer.eos_token_id is not None: + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + output = { + "prompt_input_ids": prompt_input_ids, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + if "pixel_values" in processed_features: + output["pixel_values"] = processed_features["pixel_values"][0] + if "pixel_attention_mask" in processed_features: + output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] + if "image_sizes" in processed_features: + output["image_sizes"] = processed_features["image_sizes"][0] + if "token_type_ids" in processed_features: + token_type_ids = processed_features["token_type_ids"][0] + if max_prompt_length is not None: + token_type_ids = token_type_ids[-max_prompt_length:] + output["token_type_ids"] = token_type_ids + if "pixel_position_ids" in processed_features: + output["pixel_position_ids"] = processed_features["pixel_position_ids"][0] + if "image_position_ids" in processed_features: + output["image_position_ids"] = processed_features["image_position_ids"][0] + if "mm_token_type_ids" in processed_features: + mm_token_type_ids = processed_features["mm_token_type_ids"][0] + if max_prompt_length is not None: + mm_token_type_ids = mm_token_type_ids[-max_prompt_length:] + output["mm_token_type_ids"] = mm_token_type_ids + + return output +@dataclass +class UnslothDPOConfig(DPOConfig): + """ + + Configuration class for the [`DPOTrainer`]. + + This class includes only the parameters that are specific to DPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`DPOTrainer`] is provided as a string. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to + [`~transformers.AutoModelForCausalLM.from_pretrained`] and + [`~transformers.AutoProcessor.from_pretrained`]. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + + > Parameters that control the data preprocessing + + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the left or + right depending on the `truncation_mode`. If `None`, no truncation is applied. + truncation_mode (`str`, *optional*, defaults to `"keep_start"`): + Truncation mode to use when the sequence exceeds `max_length`. The only supported value is + `"keep_start"`. The `"keep_end"` value is deprecated and will be removed in v2.0.0. + padding_free (`bool`, *optional*, defaults to `False`): + Whether to perform forward passes without padding by flattening all sequences in the batch into a single + continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only + supported with the FlashAttention 2 or 3, which can efficiently handle the flattened batch structure. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute the reference model log probabilities for the entire training dataset before + training. This allows to save memory during training, as the reference model does not need to be kept in + memory. + precompute_ref_batch_size (`int`, *optional*): + Batch size to use when precomputing reference model log probabilities. This can be set higher than the + training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for + training and `per_device_eval_batch_size` for evaluation. + + > Parameters that control the training + + loss_type (`list[str]`, *optional*, defaults to `["sigmoid"]`): + Type of loss to use. Possible values are: `'sigmoid'`, `'hinge'`, `'ipo'`, `'exo_pair'`, `'nca_pair'`, + `'robust'`, `'bco_pair'`, `'sppo_hard'`, `'aot'`, `'aot_unpaired'`, `'apo_zero'`, `'apo_down'`, + `'discopop'`, `'sft'`, `'sigmoid_norm'`. If multiple loss types are provided, they will be combined using + the weights specified in `loss_weights`. + loss_weights (`list[float]`, *optional*): + List of loss weights for multi-loss combinations. Used when combining multiple loss types. Example: `[0.8, + 0.2, 1.0]` for MPO. If not provided, defaults to equal weights (`1.0`) for all loss types. + ld_alpha (`float`, *optional*): + α parameter from the LD-DPO paper, which controls the weighting of the verbose token log-probabilities in + responses. If `None`, no weighting is applied to the verbose part, and the loss is equivalent to the + standard DPO loss. Must be in [0.0, 1.0]: `ld_alpha=1.0` applies no weighting, and `ld_alpha=0.0` masks + tokens beyond shared lengths. + f_divergence_type (`str`, *optional*, defaults to `"reverse_kl"`): + f-divergence regularizer between policy and reference (f-DPO paper). Possible values are: `reverse_kl` + (default), `forward_kl`, `js_divergence`, `alpha_divergence`. + f_alpha_divergence_coef (`float`, *optional*, defaults to `0.5`): + α coefficient for the α-divergence u^-α regularizer, used only when `f_divergence_type='alpha_divergence'`. + label_smoothing (`float`, *optional*, defaults to `0.0`): + Label smoothing parameter used in Robust DPO and EXO. In Robust DPO, it is interpreted as the probability + that a preference label is flipped and must lie in [0.0, 0.5); a typical value recommended by the Robust + DPO paper is 0.1. In EXO, it corresponds to the ε label smoothing parameter, for which the paper recommends + a typical value of 1e-3. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type='ipo'`), this value is the regularization parameter denoted + by τ in the [paper](https://huggingface.co/papers/2310.12036). + use_weighting (`bool`, *optional*, defaults to `False`): + Whether to apply WPO-style weighting (https://huggingface.co/papers/2406.11827) to preference pairs using + the policy's length-normalized sequence probabilities. + discopop_tau (`float`, *optional*, defaults to `0.05`): + τ/temperature parameter from the DiscoPOP paper, which controls the shape of the log-ratio modulated loss + when using `loss_type='discopop'`. The paper recommends the default value `discopop_tau=0.05`. + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible with + PEFT or `precompute_ref_log_probs=True`. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the TR-DPO paper, which controls the mix between the current policy and the previous + reference policy during updates. The reference policy is updated according to the equation: `π_ref = α * + π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the TR-DPO paper, which determines how frequently the current policy is synchronized with + the reference policy. To use this parameter, you must set `sync_ref_model=True`. + + > Deprecated parameters + + pad_token: + + + + Parameter `pad_token` is deprecated and will be removed in version v2.0.0. Set `tokenizer.pad_token` + directly and pass it as `processing_class` to the trainer instead. + + + + > [!NOTE] + > These parameters have default values different from [`~transformers.TrainingArguments`]: + > - `logging_steps`: Defaults to `10` instead of `500`. + > - `gradient_checkpointing`: Defaults to `True` instead of `False`. + > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. + > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. + + """ + vllm_sampling_params: Optional[Any] = field( + default = None, + metadata = {'help': 'vLLM SamplingParams'}, + ) + unsloth_num_chunks : Optional[int] = field( + default = -1, + metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, + ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {'help': 'Multiplier for chunked logit computations.'}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, + ) + max_seq_length : Optional[int] = field( + default = None, + metadata = {'help': 'Maximum sequence length to truncate to.'}, + ) + def __init__( + self, + output_dir = None, + per_device_train_batch_size = 4, + num_train_epochs = 3.0, + max_steps = -1, + learning_rate = 5e-05, + lr_scheduler_type = 'linear', + lr_scheduler_kwargs = None, + warmup_steps = 0.1, + optim = 'adamw_8bit', + optim_args = None, + weight_decay = 0.001, + adam_beta1 = 0.9, + adam_beta2 = 0.999, + adam_epsilon = 1e-08, + optim_target_modules = None, + gradient_accumulation_steps = 2, + average_tokens_across_devices = True, + max_grad_norm = 1.0, + label_smoothing_factor = 0.0, + bf16 = False, + fp16 = False, + bf16_full_eval = False, + fp16_full_eval = False, + tf32 = None, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = None, + torch_compile = False, + torch_compile_backend = None, + torch_compile_mode = None, + use_liger_kernel = False, + liger_kernel_config = None, + use_cache = False, + neftune_noise_alpha = None, + torch_empty_cache_steps = 250, + auto_find_batch_size = False, + logging_strategy = 'steps', + logging_steps = 1, + logging_first_step = False, + log_on_each_node = True, + logging_nan_inf_filter = False, + include_num_input_tokens_seen = False, + log_level = 'passive', + log_level_replica = 'warning', + disable_tqdm = None, + report_to = 'none', + run_name = None, + project = 'huggingface', + trackio_space_id = 'trackio', + eval_strategy = 'no', + eval_steps = None, + eval_delay = 0, + per_device_eval_batch_size = 4, + prediction_loss_only = False, + eval_on_start = False, + eval_do_concat_batches = True, + eval_use_gather_object = False, + eval_accumulation_steps = 2, + batch_eval_metrics = False, + save_only_model = False, + save_strategy = 'steps', + save_steps = 500, + save_on_each_node = False, + save_total_limit = None, + enable_jit_checkpoint = False, + push_to_hub = False, + hub_token = None, + hub_private_repo = None, + hub_model_id = None, + hub_strategy = 'every_save', + hub_always_push = False, + hub_revision = None, + load_best_model_at_end = False, + metric_for_best_model = None, + greater_is_better = None, + ignore_data_skip = False, + restore_callback_states_from_checkpoint = False, + full_determinism = False, + seed = 3407, + data_seed = 3407, + use_cpu = False, + accelerator_config = None, + parallelism_config = None, + dataloader_drop_last = False, + dataloader_num_workers = 0, + dataloader_pin_memory = True, + dataloader_persistent_workers = False, + dataloader_prefetch_factor = None, + remove_unused_columns = True, + label_names = None, + train_sampling_strategy = 'random', + length_column_name = 'length', + ddp_find_unused_parameters = None, + ddp_bucket_cap_mb = None, + ddp_broadcast_buffers = None, + ddp_backend = None, + ddp_timeout = 1800, + fsdp = None, + fsdp_config = None, + deepspeed = None, + debug = '', + skip_memory_metrics = True, + do_train = False, + do_eval = False, + do_predict = False, + resume_from_checkpoint = None, + warmup_ratio = None, + logging_dir = None, + local_rank = -1, + model_init_kwargs = None, + trust_remote_code = False, + disable_dropout = True, + dataset_num_proc = None, + max_length = 1024, + truncation_mode = 'keep_start', + padding_free = None, + pad_to_multiple_of = None, + precompute_ref_log_probs = False, + precompute_ref_batch_size = None, + loss_weights = None, + ld_alpha = None, + f_divergence_type = 'reverse_kl', + f_alpha_divergence_coef = 0.5, + label_smoothing = 0.0, + beta = 0.1, + use_weighting = False, + discopop_tau = 0.05, + activation_offloading = False, + sync_ref_model = False, + ref_model_mixup_alpha = 0.6, + ref_model_sync_steps = 512, + pad_token = None, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, + max_seq_length = None, + **kwargs, + ): + if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') + if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') + if num_train_epochs is None: + num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override + if output_dir is None and save_strategy == 'steps' and save_steps == 500: + output_dir = 'unsloth_training_checkpoints' + save_strategy = 'no' + import multiprocessing as _mp + if dataset_num_proc is None: + if _mp.get_start_method() != 'fork': + dataset_num_proc = None + else: + import psutil + dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64) + memory_gb_left = psutil.virtual_memory().available / (1024**3) + if memory_gb_left <= 2: dataset_num_proc = 1 + else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left)) + if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': + from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION + if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: + from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE + pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE + + + super().__init__( + output_dir = output_dir, + per_device_train_batch_size = per_device_train_batch_size, + num_train_epochs = num_train_epochs, + max_steps = max_steps, + learning_rate = learning_rate, + lr_scheduler_type = lr_scheduler_type, + lr_scheduler_kwargs = lr_scheduler_kwargs, + warmup_steps = warmup_steps, + optim = optim, + optim_args = optim_args, + weight_decay = weight_decay, + adam_beta1 = adam_beta1, + adam_beta2 = adam_beta2, + adam_epsilon = adam_epsilon, + optim_target_modules = optim_target_modules, + gradient_accumulation_steps = gradient_accumulation_steps, + average_tokens_across_devices = average_tokens_across_devices, + max_grad_norm = max_grad_norm, + label_smoothing_factor = label_smoothing_factor, + bf16 = bf16, + fp16 = fp16, + bf16_full_eval = bf16_full_eval, + fp16_full_eval = fp16_full_eval, + tf32 = tf32, + gradient_checkpointing = gradient_checkpointing, + gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, + torch_compile = torch_compile, + torch_compile_backend = torch_compile_backend, + torch_compile_mode = torch_compile_mode, + use_liger_kernel = use_liger_kernel, + liger_kernel_config = liger_kernel_config, + use_cache = use_cache, + neftune_noise_alpha = neftune_noise_alpha, + torch_empty_cache_steps = torch_empty_cache_steps, + auto_find_batch_size = auto_find_batch_size, + logging_strategy = logging_strategy, + logging_steps = logging_steps, + logging_first_step = logging_first_step, + log_on_each_node = log_on_each_node, + logging_nan_inf_filter = logging_nan_inf_filter, + include_num_input_tokens_seen = include_num_input_tokens_seen, + log_level = log_level, + log_level_replica = log_level_replica, + disable_tqdm = disable_tqdm, + report_to = report_to, + run_name = run_name, + project = project, + trackio_space_id = trackio_space_id, + eval_strategy = eval_strategy, + eval_steps = eval_steps, + eval_delay = eval_delay, + per_device_eval_batch_size = per_device_eval_batch_size, + prediction_loss_only = prediction_loss_only, + eval_on_start = eval_on_start, + eval_do_concat_batches = eval_do_concat_batches, + eval_use_gather_object = eval_use_gather_object, + eval_accumulation_steps = eval_accumulation_steps, + batch_eval_metrics = batch_eval_metrics, + save_only_model = save_only_model, + save_strategy = save_strategy, + save_steps = save_steps, + save_on_each_node = save_on_each_node, + save_total_limit = save_total_limit, + enable_jit_checkpoint = enable_jit_checkpoint, + push_to_hub = push_to_hub, + hub_token = hub_token, + hub_private_repo = hub_private_repo, + hub_model_id = hub_model_id, + hub_strategy = hub_strategy, + hub_always_push = hub_always_push, + hub_revision = hub_revision, + load_best_model_at_end = load_best_model_at_end, + metric_for_best_model = metric_for_best_model, + greater_is_better = greater_is_better, + ignore_data_skip = ignore_data_skip, + restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, + full_determinism = full_determinism, + seed = seed, + data_seed = data_seed, + use_cpu = use_cpu, + accelerator_config = accelerator_config, + parallelism_config = parallelism_config, + dataloader_drop_last = dataloader_drop_last, + dataloader_num_workers = dataloader_num_workers, + dataloader_pin_memory = dataloader_pin_memory, + dataloader_persistent_workers = dataloader_persistent_workers, + dataloader_prefetch_factor = dataloader_prefetch_factor, + remove_unused_columns = remove_unused_columns, + label_names = label_names, + train_sampling_strategy = train_sampling_strategy, + length_column_name = length_column_name, + ddp_find_unused_parameters = ddp_find_unused_parameters, + ddp_bucket_cap_mb = ddp_bucket_cap_mb, + ddp_broadcast_buffers = ddp_broadcast_buffers, + ddp_backend = ddp_backend, + ddp_timeout = ddp_timeout, + fsdp = fsdp, + fsdp_config = fsdp_config, + deepspeed = deepspeed, + debug = debug, + skip_memory_metrics = skip_memory_metrics, + do_train = do_train, + do_eval = do_eval, + do_predict = do_predict, + resume_from_checkpoint = resume_from_checkpoint, + warmup_ratio = warmup_ratio, + logging_dir = logging_dir, + local_rank = local_rank, + model_init_kwargs = model_init_kwargs, + trust_remote_code = trust_remote_code, + disable_dropout = disable_dropout, + dataset_num_proc = dataset_num_proc, + max_length = max_length, + truncation_mode = truncation_mode, + padding_free = padding_free, + pad_to_multiple_of = pad_to_multiple_of, + precompute_ref_log_probs = precompute_ref_log_probs, + precompute_ref_batch_size = precompute_ref_batch_size, + loss_weights = loss_weights, + ld_alpha = ld_alpha, + f_divergence_type = f_divergence_type, + f_alpha_divergence_coef = f_alpha_divergence_coef, + label_smoothing = label_smoothing, + beta = beta, + use_weighting = use_weighting, + discopop_tau = discopop_tau, + activation_offloading = activation_offloading, + sync_ref_model = sync_ref_model, + ref_model_mixup_alpha = ref_model_mixup_alpha, + ref_model_sync_steps = ref_model_sync_steps, + pad_token = pad_token,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier + self.max_seq_length = max_seq_length + # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+ + if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: + if 'use_reentrant' in self.gradient_checkpointing_kwargs: + del self.gradient_checkpointing_kwargs['use_reentrant'] + +pass + +class _UnslothDPOTrainer(_BaseTrainer): + """""" + + _tag_names = ["trl", "dpo"] + _name = "DPO" + _paper = { + "title": "Direct Preference Optimization: Your Language Model is Secretly a Reward Model", + "id": "2305.18290", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{rafailov2023direct, + title = {{Direct Preference Optimization: Your Language Model is Secretly a Reward Model}}, + author = {Rafael Rafailov and Archit Sharma and Eric Mitchell and Christopher D. Manning and Stefano Ermon and Chelsea Finn}, + year = 2023, + booktitle = {Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023}, + url = {http://papers.nips.cc/paper_files/paper/2023/hash/a85b405ed65c6477a4fe8302b5e06ce7-Abstract-Conference.html}, + editor = {Alice Oh and Tristan Naumann and Amir Globerson and Kate Saenko and Moritz Hardt and Sergey Levine}, + }"""), + } + + def __init__( + self, + model: "str | PreTrainedModel | PeftModel", + ref_model: PreTrainedModel | None = None, + args: DPOConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + peft_config: "PeftConfig | None" = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = DPOConfig(f"{model_name}-DPO") + + if train_dataset is None: + raise ValueError("`train_dataset` is required") + elif isinstance(train_dataset, IterableDataset): + # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate + # batches from multiple processes, leading to mismatch errors. + if args.accelerator_config.dispatch_batches is True: + logger.warning( + "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " + "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " + "`dispatch_batches` in `DPOConfig` or set it to `False`." + ) + args.accelerator_config.dispatch_batches = False + + # Model + if isinstance(model, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + model = create_model_from_path(model, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `DPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do + _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. In most cases you should omit `ref_model` and " + "we'll initialize it to a copy of `model` for you." + ) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained( + get_config_model_id(model.config), trust_remote_code=args.trust_remote_code + ) + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + self._tokenizer = processing_class.tokenizer + self._is_vlm = True + elif isinstance(processing_class, PreTrainedTokenizerBase): + self._tokenizer = processing_class + self._is_vlm = False + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if self._tokenizer.pad_token is None: + self._tokenizer.pad_token = self._tokenizer.eos_token + + # PEFT + if False: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if is_peft_model(model): + raise ValueError( + "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " + "and unload the existing adapter, save the resulting base model, and then pass that base model along " + "with the new `peft_config` to the trainer." + ) + # Create PEFT model + # ZeRO-3 + PEFT for non-quantized models: + # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. + # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, + # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. + # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. + # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for + # non-quantized bf16 training. + # - See: + # - TRL issue: https://github.com/huggingface/trl/issues/6089 + # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 + # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg + get_peft_model_kwargs = {} + if ( + args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and not _is_quantized_model + and Version(peft.__version__) >= Version("0.12.0") + ): + get_peft_model_kwargs["autocast_adapter_dtype"] = False + model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + + elif is_peft_model(model) and ref_model is None: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during DPO training. PEFT only + # supports one adapter per model when the LoRA config uses `target_parameters` [see peft#3340], so in that + # case we skip the "ref" adapter and compute the reference log probs with adapters disabled, i.e. with the + # base model. + default_config = model.peft_config["default"] + if isinstance(default_config, LoraConfig) and default_config.target_parameters: + logger.warning( + "PEFT can't add a frozen reference adapter alongside one that uses `target_parameters` " + "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " + "If you wrapped the model only to apply LoRA, pass a `peft_config` to the trainer instead; if you " + "wrapped it deliberately (pretrained adapter or custom init), note that the base model matches " + "your adapter only when it's freshly zero-initialized. If it is, this warning is safe to ignore." + ) + else: + model.add_adapter("ref", default_config) + for name, param in model.named_parameters(): + if ".default." in name: + ref_name = name.replace(".default.", ".ref.") + ref_param = model.get_parameter(ref_name) + ref_param.data.copy_(param.data) + + # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally + # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 + if is_peft_model(model) and args.gradient_checkpointing: + model.enable_input_require_grads() + + # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the + # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by + # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for + # quantized models. See: https://github.com/huggingface/peft/issues/2889 + if _is_quantized_model: + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.bfloat16) + + # Data collator + self.padding_free = args.padding_free + if self.padding_free: + logger.warning( + "`padding_free=True` is temporarily unavailable after a refactor and is currently disabled. Falling " + "back to standard padding (`padding_free=False`). This feature is planned to return in a future " + "update; for now, please set `padding_free=False` explicitly." + ) + self.padding_free = False + dataset_sample = next(iter(train_dataset)) + self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample + if self._is_vision_dataset and not self._is_vlm: + raise ValueError( + "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " + "model does not seem to be a vision-language model. Please check your model and dataset." + ) + if self._is_vision_dataset and args.max_length is not None and args.truncation_mode == "keep_end": + raise ValueError( + "truncation_mode='keep_end' is not supported for vision-language models. Image tokens reside " + "inside the prompt portion of the sequence; depending on the example, keep_end may silently " + "drop them, causing pixel_values to be forwarded to the model with no corresponding visual " + "tokens in input_ids. Use truncation_mode='keep_start' (the default) or set max_length=None." + ) + if self._is_vision_dataset and args.precompute_ref_log_probs: + raise ValueError( + "`precompute_ref_log_probs=True` is not supported for vision datasets. For vision-language " + "models, all data processing is performed on the fly rather than upfront, and running a full " + "forward pass of the reference model over the entire dataset is not supported for large " + "multimodal models. Set `precompute_ref_log_probs=False`." + ) + if data_collator is None and not self._is_vision_dataset: + # Get the pad token: if not provided, use the one from the processing class or the eos token + # if the processing class does not have a pad token. + pad_token = args.pad_token or self._tokenizer.pad_token or self._tokenizer.eos_token + if pad_token not in self._tokenizer.get_vocab(): + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + self._tokenizer.pad_token = pad_token + data_collator = DataCollatorForPreference( + pad_token_id=self._tokenizer.pad_token_id, + max_length=args.max_length, + truncation_mode=args.truncation_mode, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + elif data_collator is None and self._is_vision_dataset: + data_collator = DataCollatorForVisionPreference( + processor=processing_class, + max_length=args.max_length, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + + # Training arguments + self.beta = args.beta + self.precompute_ref_logps = args.precompute_ref_log_probs + self.loss_types = args.loss_type # args.loss_type is already a list + self.loss_weights = args.loss_weights or [1.0] * len(self.loss_types) + self.ld_alpha = args.ld_alpha + self.f_divergence_type = args.f_divergence_type + self.f_alpha_divergence_coef = args.f_alpha_divergence_coef + self.label_smoothing = args.label_smoothing + self.use_weighting = args.use_weighting + if self.use_weighting and any(loss_type in {"aot", "aot_unpaired"} for loss_type in self.loss_types): + raise NotImplementedError( + "WPO-style weighting is not implemented for 'aot' or 'aot_unpaired' because those losses sort " + "samples, which would misalign per-pair weights." + ) + if "robust" in self.loss_types and not (0.0 <= self.label_smoothing < 0.5): + logger.warning( + "The `label_smoothing` parameter should lie in [0.0, 0.5) for the 'robust' loss. You provided " + f"{self.label_smoothing}." + ) + if "exo_pair" in self.loss_types and self.label_smoothing == 0.0: + raise ValueError( + "Label smoothing must be greater than 0.0 when using 'exo_pair' loss. The EXO paper recommends a " + "value of 1e-3." + ) + self.use_liger_kernel = args.use_liger_kernel + if args.use_liger_kernel: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_kernel=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if len(self.loss_types) != 1: + raise NotImplementedError( + "Multiple loss types are not yet supported when using Liger kernel. If you need this feature, " + "please open a feature request at https://github.com/huggingface/trl/issues." + ) + self.liger_loss_fn = LigerFusedLinearDPOLoss(beta=args.beta, loss_type=self.loss_types[0]) + if compute_metrics is not None: + raise ValueError( + "compute_metrics is not supported with the Liger kernel. compute_metrics requires to be able to " + "recover the logits from the forward pass, but Liger kernel does not materialize logits." + ) + if self.precompute_ref_logps: + raise ValueError( + "Liger DPO loss does not support precomputing reference log probabilities. Either disable " + "`precompute_ref_log_probs` or set `use_liger_kernel` to False." + ) + if is_peft_model(model): + raise NotImplementedError("Liger DPO loss is not implemented for PEFT models.") + + # Dataset + # Skip dataset preparation if it's a VLM, where preprocessing [e.g., image-to-pixel conversion] is too costly + # and done on the fly instead. + skip_prepare_dataset = self._is_vision_dataset + if not skip_prepare_dataset: + train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + ) + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # Reference model + if ref_model is None: + if is_peft_model(self.model) or args.precompute_ref_log_probs: + # If PEFT is used, the reference model is not needed since the adapter can be disabled to revert to the + # initial model. If precompute_ref_log_probs is True, the reference model does not need to be kept in + # memory during training. + self.ref_model = None + else: + ref_model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + ref_model_init_kwargs["device_map"] = None + ref_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + ref_model_path = get_config_model_id(self.model.config) + self.ref_model = create_model_from_path(ref_model_path, **ref_model_init_kwargs) + else: + self.ref_model = ref_model + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + if is_peft_model(self.model): + raise NotImplementedError( + "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " + "With PEFT, DPOTrainer does not keep a separate reference model in memory; instead, it recovers " + "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " + "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " + "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " + "feature request at https://github.com/huggingface/trl/issues." + ) + if args.precompute_ref_log_probs: + raise ValueError( + "You cannot use `sync_ref_model=True` together with `precompute_ref_log_probs=True`. " + "`precompute_ref_log_probs=True` assumes a fixed reference model, but with `sync_ref_model=True` " + "the reference model is periodically updated during training, making any precomputed reference " + "log-probs stale. Set `precompute_ref_log_probs=False` or disable `sync_ref_model`." + ) + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + if args.precompute_ref_log_probs: + if isinstance(self.train_dataset, IterableDataset) or isinstance( + self.eval_dataset, (IterableDataset, IterableDatasetDict) + ): + raise ValueError( + "`precompute_ref_log_probs=True` is not supported with IterableDataset. Please use a map-style " + "Dataset or set `precompute_ref_log_probs=False`." + ) + + self.train_dataset = self._precompute_ref_logps( + self.train_dataset, + "train", + self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size, + ) + if self.eval_dataset is not None: + if isinstance(self.eval_dataset, dict): + self.eval_dataset = { + name: self._precompute_ref_logps( + dataset, name, self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size + ) + for name, dataset in self.eval_dataset.items() + } + else: + self.eval_dataset = self._precompute_ref_logps( + self.eval_dataset, + "eval", + self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size, + ) + + def _tokenize( + self, + processing_class: PreTrainedTokenizerBase | ProcessorMixin, + input: str | list, + **kwargs, + ) -> dict[str, list]: + """Tokenize a single example for dataset preprocessing. + + Dispatches to `apply_chat_template` for conversational input (list of message dicts) and to `__call__` for + non-conversational input (str). For VLMs, normalizes the batch dimension that processors emit even for single + examples. + + Args: + processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`]): + The tokenizer or processor to use. + input (`str` or `list`): + A string for non-conversational input, or a list of message dicts for conversational input. + **kwargs: + Forwarded to `apply_chat_template` (e.g. `add_generation_prompt`, `return_assistant_tokens_mask`). + + Returns: + `dict` with at least an `"input_ids"` key mapping to a flat `list[int]`. + """ + if isinstance(input, list): # conversational: list of message dicts + if self._is_vlm: + input = prepare_multimodal_messages(input) + result = processing_class.apply_chat_template(input, tokenize=True, return_dict=True, **kwargs) + else: # non-conversational: plain text string + result = processing_class(text=input) + # VLMs emit a batch dimension even for single examples; unwrap it + if self._is_vlm: + return {k: v[0] for k, v in result.items()} + return result + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase | ProcessorMixin, + args: DPOConfig, + dataset_name: str, + ) -> Dataset | IterableDataset: + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + with PartialState().main_process_first(): + # Extract the prompt if needed + first_example = next(iter(dataset)) + if "prompt" not in first_example: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Extracting prompt from {dataset_name} dataset" + dataset = dataset.map(extract_prompt, **map_kwargs) + + # Add EOS token if needed: non-conversational only + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if not example["chosen"].endswith(eos_token): + example["chosen"] = example["chosen"] + eos_token + if not example["rejected"].endswith(eos_token): + example["rejected"] = example["rejected"] + eos_token + return example + + dataset = dataset.map(add_eos, fn_kwargs={"eos_token": self._tokenizer.eos_token}, **map_kwargs) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + def tokenize_fn(example, processing_class): + tools = example.get("tools") + tools = json.loads(tools) if isinstance(tools, str) else tools + output = {} + if is_conversational(example): + prompt_ids = self._tokenize( + processing_class, + example["prompt"], + tools=tools, + add_generation_prompt=True, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + prompt_chosen_ids = self._tokenize( + processing_class, + example["prompt"] + example["chosen"], + tools=tools, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + prompt_rejected_ids = self._tokenize( + processing_class, + example["prompt"] + example["rejected"], + tools=tools, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + else: + prompt_ids = self._tokenize(processing_class, example["prompt"])["input_ids"] + prompt_chosen_ids = self._tokenize(processing_class, example["prompt"] + example["chosen"])[ + "input_ids" + ] + prompt_rejected_ids = self._tokenize(processing_class, example["prompt"] + example["rejected"])[ + "input_ids" + ] + + # Check if the tokenized prompt starts with the tokenized prompt+completion + if not prompt_chosen_ids[: len(prompt_ids)] == prompt_ids: + logger.warning( + "Mismatch between tokenized prompt and the start of tokenized prompt+chosen. " + "This may be due to unexpected tokenizer behavior, whitespace issues, or special " + "token handling. Verify that the tokenizer is processing text consistently." + ) + if not prompt_rejected_ids[: len(prompt_ids)] == prompt_ids: + logger.warning( + "Mismatch between tokenized prompt and the start of tokenized prompt+rejected. " + "This may be due to unexpected tokenizer behavior, whitespace issues, or special " + "token handling. Verify that the tokenizer is processing text consistently." + ) + + output["prompt_ids"] = prompt_ids + output["chosen_ids"] = prompt_chosen_ids[len(prompt_ids) :] + output["rejected_ids"] = prompt_rejected_ids[len(prompt_ids) :] + return output + + dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) + + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). + if self._signature_columns is None: + if self._is_vision_dataset: + self._signature_columns = [ + "prompt", + "chosen", + "rejected", + "image", + "images", + "tools", + "chat_template_kwargs", + ] + else: + self._signature_columns = [ + "prompt_ids", + "chosen_ids", + "rejected_ids", + "ref_chosen_logps", + "ref_rejected_logps", + ] + + def _precompute_ref_logps(self, dataset: Dataset, name: str, batch_size: int) -> Dataset: + model_hash = hash_module(self.ref_model or self.model) + fingerprint = Hasher.hash((dataset._fingerprint, model_hash)) + cache_file = dataset._get_cache_file_path(fingerprint) + + if os.path.exists(cache_file): + return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) + + dataloader = DataLoader( + dataset, + batch_size=batch_size, + collate_fn=self.data_collator, + num_workers=self.args.dataloader_num_workers, + pin_memory=self.args.dataloader_pin_memory, + shuffle=False, + ) + data_loader = self.accelerator.prepare(dataloader) + ref_chosen_logps = [] + ref_rejected_logps = [] + for padded_batch in tqdm(iterable=data_loader, desc=f"Computing reference log probs for {name} dataset"): + ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch) + ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics( + (ref_chosen_logp, ref_rejected_logp) + ) + ref_chosen_logps.append(ref_chosen_logp.cpu()) + ref_rejected_logps.append(ref_rejected_logp.cpu()) + + ref_chosen_logps = torch.cat(ref_chosen_logps) + ref_rejected_logps = torch.cat(ref_rejected_logps) + + if self.accelerator.is_main_process: + + def add_ref_logps(batch, indices): + return { + "ref_chosen_logps": ref_chosen_logps[indices], + "ref_rejected_logps": ref_rejected_logps[indices], + } + + dataset.map( + add_ref_logps, + with_indices=True, + batched=True, + remove_columns=dataset.column_names, + new_fingerprint=fingerprint, + desc=f"Caching reference log probs for {name} dataset", + ) + self.accelerator.wait_for_everyone() + + return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) + + def compute_ref_log_probs(self, inputs): + """Computes reference log probabilities for a single padded batch.""" + device = self.accelerator.device + + _non_model_keys = {"completion_mask", "ref_chosen_logps", "ref_rejected_logps"} + model_kwargs = {k: v for k, v in inputs.items() if k not in _non_model_keys} + model_kwargs["use_cache"] = False + + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + if self.ref_model is None: + if is_peft_model(self.model): + model = self.accelerator.unwrap_model(self.model) + with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): + ref_outputs = self.model(**model_kwargs) + else: + ref_outputs = self.model(**model_kwargs) + else: + ref_outputs = self.ref_model(**model_kwargs) + + input_ids = inputs["input_ids"] + completion_mask = inputs["completion_mask"] + shift_labels = input_ids[..., 1:] + shift_completion_mask = completion_mask[..., 1:] + ref_shift_logits = ref_outputs.logits[..., :-1, :] + ref_per_token_logps = selective_log_softmax(ref_shift_logits, shift_labels) + ref_per_token_logps[shift_completion_mask == 0] = 0.0 + + if self.ld_alpha is None: + ref_logps = ref_per_token_logps.sum(dim=1) + else: + comp_pos = shift_completion_mask.cumsum(dim=1) + comp_lens = shift_completion_mask.sum(dim=1).long() + chosen_lens, rejected_lens = comp_lens.chunk(2, dim=0) + shared_lens = torch.minimum(chosen_lens, rejected_lens) + shared_lens = torch.cat([shared_lens, shared_lens], dim=0).to(device) + shared_mask = (comp_pos > 0) & (comp_pos <= shared_lens.unsqueeze(1)) + tail_mask = comp_pos > shared_lens.unsqueeze(1) + shared_logps = (ref_per_token_logps * shared_mask).sum(dim=1) + tail_logps = (ref_per_token_logps * tail_mask).sum(dim=1) + ref_logps = shared_logps + self.ld_alpha * tail_logps + + ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0) + return ref_chosen_logps, ref_rejected_logps + + def _compute_loss_liger(self, model, inputs, return_outputs): + if return_outputs: + raise RuntimeError( + "return_outputs=True is not supported with the Liger DPO loss. The Liger loss computes the loss " + "without materializing logits, so outputs cannot be returned." + ) + + mode = "train" if self.model.training else "eval" + + _non_model_keys = {"completion_mask", "ref_chosen_logps", "ref_rejected_logps"} + model_kwargs = {k: v for k, v in inputs.items() if k not in _non_model_keys} + model_kwargs["use_cache"] = False + + # `base_model` gives the backbone model (skipping `lm_head`) — text decoder for LMs, multimodal wrapper for + # VLMs (so vision-token injection runs before the text decoder). `get_decoder()` won't do: on VLMs it + # returns just the text stack and feeds image-placeholder IDs through it. + # Pre-5.0 transformers VLMs set `base_model_prefix = ""` so `base_model is self` (re-runs `lm_head`). + # Fall back to `.model` there. + if self._is_vlm and Version(transformers.__version__) < Version("5.0.0"): + backbone, ref_backbone = model.model, self.ref_model.model + else: + backbone, ref_backbone = model.base_model, self.ref_model.base_model + + outputs = backbone(**model_kwargs) + hidden_states = outputs.last_hidden_state[:, :-1].contiguous() + lm_head = model.get_output_embeddings() + weight = lm_head.weight + bias = lm_head.bias + + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + ref_outputs = ref_backbone(**model_kwargs) + ref_lm_head = self.ref_model.get_output_embeddings() + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1].contiguous() + ref_weight = ref_lm_head.weight + ref_bias = ref_lm_head.bias + + input_ids = model_kwargs["input_ids"] + completion_mask = inputs["completion_mask"] + shift_completion_mask = completion_mask[:, 1:] + labels = input_ids[:, 1:].clone() + labels[shift_completion_mask == 0] = -100 + + loss, metrics = self.liger_loss_fn( + weight, hidden_states, labels, bias, ref_hidden_states, ref_weight, ref_bias + ) + + ( + chosen_logps, + rejected_logps, + chosen_logits_mean, + rejected_logits_mean, + nll_loss, + chosen_rewards, + rejected_rewards, + ) = metrics + + if mode == "train": + num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + self._total_train_tokens += num_tokens_in_batch + self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + + avg_chosen_logits = self.accelerator.gather_for_metrics(chosen_logits_mean).mean().item() + avg_rejected_logits = self.accelerator.gather_for_metrics(rejected_logits_mean).mean().item() + self._metrics[mode]["logits/chosen"].append(avg_chosen_logits) + self._metrics[mode]["logits/rejected"].append(avg_rejected_logits) + + agg_chosen_rewards = self.accelerator.gather(chosen_rewards) + agg_rejected_rewards = self.accelerator.gather(rejected_rewards) + self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item()) + self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item()) + + reward_accuracies = (chosen_rewards > rejected_rewards).float() + agg_reward_accuracies = self.accelerator.gather(reward_accuracies) + self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item()) + + margins = chosen_rewards - rejected_rewards + agg_margins = self.accelerator.gather(margins) + self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item()) + + self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item()) + self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item()) + + return loss + + def _compute_loss(self, model, inputs, return_outputs): + mode = "train" if self.model.training else "eval" + device = self.accelerator.device + + _non_model_keys = {"completion_mask", "ref_chosen_logps", "ref_rejected_logps"} + model_kwargs = {k: v for k, v in inputs.items() if k not in _non_model_keys} + model_kwargs["use_cache"] = False + outputs = model(**model_kwargs) + + input_ids = inputs["input_ids"] + completion_mask = inputs["completion_mask"] + shift_logits = outputs.logits[..., :-1, :] + shift_labels = input_ids[..., 1:] + shift_completion_mask = completion_mask[..., 1:] + per_token_logps = selective_log_softmax(shift_logits, shift_labels) + per_token_logps[shift_completion_mask == 0] = 0.0 # mask out non-completion tokens + if self.ld_alpha is None: + logps = per_token_logps.sum(dim=1) # sum over sequence length + else: + comp_pos = shift_completion_mask.cumsum(dim=1) + comp_lens = shift_completion_mask.sum(dim=1).long() + chosen_lens, rejected_lens = comp_lens.chunk(2, dim=0) + shared_lens = torch.minimum(chosen_lens, rejected_lens) + shared_lens = torch.cat([shared_lens, shared_lens], dim=0).to(device) + shared_mask = (comp_pos > 0) & (comp_pos <= shared_lens.unsqueeze(1)) # shared: 1 <= pos <= shared_len + tail_mask = comp_pos > shared_lens.unsqueeze(1) # tail: pos > shared_len + shared_logps = (per_token_logps * shared_mask).sum(dim=1) + tail_logps = (per_token_logps * tail_mask).sum(dim=1) + logps = shared_logps + self.ld_alpha * tail_logps + chosen_logps, rejected_logps = logps.chunk(2, dim=0) # batch is [chosen, rejected] + + if self.precompute_ref_logps: + ref_chosen_logps, ref_rejected_logps = inputs["ref_chosen_logps"], inputs["ref_rejected_logps"] + else: + # When gradient checkpointing is enabled with use_reentrant=True (default), calling the model inside a + # torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True"). + # Temporarily disable checkpointing to avoid this warning during inference. + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + if is_peft_model(model) and self.ref_model is None: + # When training a PEFT adapter, how we obtain the reference depends on the setup: + # - New adapter: disabling adapters yields the base model. + # - Re-training an existing adapter: an initial copy is loaded under the name "ref". + model = self.accelerator.unwrap_model(model) + with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): + ref_outputs = self.model(**model_kwargs) + else: + ref_outputs = self.ref_model(**model_kwargs) + + ref_shift_logits = ref_outputs.logits[..., :-1, :] + ref_per_token_logps = selective_log_softmax(ref_shift_logits, shift_labels) + ref_per_token_logps[shift_completion_mask == 0] = 0.0 # mask out non-completion tokens + if self.ld_alpha is None: + ref_logps = ref_per_token_logps.sum(dim=1) # sum over sequence length + else: + # reuse comp_pos/shared_mask/tail_mask computed above (they depend only on completion_mask) + ref_shared_logps = (ref_per_token_logps * shared_mask).sum(dim=1) + ref_tail_logps = (ref_per_token_logps * tail_mask).sum(dim=1) + ref_logps = ref_shared_logps + self.ld_alpha * ref_tail_logps + ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0) # batch is [chosen, rejected] + + # Get the log ratios for the chosen and rejected responses + chosen_logratios = chosen_logps - ref_chosen_logps + rejected_logratios = rejected_logps - ref_rejected_logps + + if self.f_divergence_type == "reverse_kl": # standard DPO + chosen_scores = chosen_logratios + rejected_scores = rejected_logratios + elif self.f_divergence_type == "forward_kl": + # f'(t) = 1 - 1/t -> drop constant -> -exp(-logratio) + chosen_scores = -torch.exp(-chosen_logratios) + rejected_scores = -torch.exp(-rejected_logratios) + elif self.f_divergence_type == "js_divergence": + # f'(t) = log(2t/(t+1)) -> drop log 2 + chosen_scores = F.logsigmoid(chosen_logratios) + rejected_scores = F.logsigmoid(rejected_logratios) + elif self.f_divergence_type == "alpha_divergence": + # alpha-divergence: f'(t) = (t^(α-1) - 1)/(α-1) + if abs(self.f_alpha_divergence_coef - 1.0) < 1e-6: # limit case f'(t) -> log(t), fall back to reverse_kl + chosen_scores = chosen_logratios + rejected_scores = rejected_logratios + else: + coef = 1.0 / (self.f_alpha_divergence_coef - 1.0) + t_chosen = (self.f_alpha_divergence_coef - 1.0) * chosen_logratios + t_rejected = (self.f_alpha_divergence_coef - 1.0) * rejected_logratios + dtype = t_chosen.dtype + # Clamp max so exp(.) stays representable after casting back + clamp_max = {torch.float16: 11.0, torch.bfloat16: 80.0, torch.float32: 80.0}[dtype] + t_chosen_float = torch.clamp(t_chosen.float(), max=clamp_max) + t_rejected_float = torch.clamp(t_rejected.float(), max=clamp_max) + chosen_scores = torch.exp(t_chosen_float).to(dtype) * coef + rejected_scores = torch.exp(t_rejected_float).to(dtype) * coef + else: + raise ValueError(f"Unknown f_divergence_type: {self.f_divergence_type}") + + delta_score = chosen_scores - rejected_scores + + loss = 0.0 + for loss_type, loss_weight in zip(self.loss_types, self.loss_weights, strict=True): + if loss_type == "sigmoid": + per_sequence_loss = -F.logsigmoid(self.beta * delta_score) + + elif loss_type == "hinge": + per_sequence_loss = torch.relu(1 - self.beta * delta_score) + + elif loss_type == "ipo": + # IPO uses sequence-level log-prob differences; in code these are token-summed over the completion, + # which makes the squared loss scale with completion length. We therefore normalize by the number of + # completion tokens (average per token) to make β/loss comparable across variable lengths. This length + # normalization is not explicitly discussed in the IPO paper; we confirmed this choice with the IPO + # authors, and the results reported in the paper correspond to this normalized form. + chosen_mask, rejected_mask = completion_mask.chunk(2, dim=0) + chosen_avg_score = chosen_scores / chosen_mask.sum(dim=1).clamp(min=1.0) + rejected_avg_score = rejected_scores / rejected_mask.sum(dim=1).clamp(min=1.0) + ipo_delta = chosen_avg_score - rejected_avg_score + # (Eq. 17) of the paper where beta is the regularization parameter for the IPO loss, denoted by τ. + per_sequence_loss = (ipo_delta - 1 / (2 * self.beta)) ** 2 + + elif loss_type == "exo_pair": + # Implements EXO-pref from the paper https://huggingface.co/papers/2402.00856, (Eq. 16) + # Minimize KL(p_fθ || p_rh) for K=2; p_fθ = softmax(βπ * (log πθ − log π_ref)) over {chosen, rejected} + # p_rh = [(1−ε), ε]; expanded KL gives the weighted logsigmoid form below + epsilon = torch.tensor(self.label_smoothing, device=device) + qw = torch.sigmoid(self.beta * delta_score) + log_qw = F.logsigmoid(self.beta * delta_score) + log_pw = torch.log1p(-epsilon) + ql = torch.sigmoid(-self.beta * delta_score) + log_ql = F.logsigmoid(-self.beta * delta_score) + log_pl = torch.log(epsilon) + per_sequence_loss = qw * (log_qw - log_pw) + ql * (log_ql - log_pl) + + elif loss_type == "nca_pair": + chosen_rewards = self.beta * chosen_scores + rejected_rewards = self.beta * rejected_scores + per_sequence_loss = ( + -F.logsigmoid(chosen_rewards) + - 0.5 * F.logsigmoid(-chosen_rewards) + - 0.5 * F.logsigmoid(-rejected_rewards) + ) + + elif loss_type == "robust": + clean_loss_term = -(1 - self.label_smoothing) * F.logsigmoid(self.beta * delta_score) + flipped_loss_term = -self.label_smoothing * F.logsigmoid(-self.beta * delta_score) + per_sequence_loss = (clean_loss_term - flipped_loss_term) / (1 - 2 * self.label_smoothing) + + elif loss_type == "bco_pair": + chosen_rewards = self.beta * chosen_scores + rejected_rewards = self.beta * rejected_scores + per_sequence_loss = -F.logsigmoid(chosen_rewards) - F.logsigmoid(-rejected_rewards) + + elif loss_type == "sppo_hard": + # In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach, + # estimated using the PairRM score. The probability calculation is conducted outside of the trainer + # class. The version described here is the hard probability version, where P in Equation (4.7) of + # Algorithm 1 is set to 1 for the winner and 0 for the loser. + winner_margin_error = (chosen_scores - 0.5 / self.beta) ** 2 + loser_margin_error = (rejected_scores + 0.5 / self.beta) ** 2 + per_sequence_loss = winner_margin_error + loser_margin_error + + elif loss_type == "aot": + logratios = chosen_logps - rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logratios_sorted, _ = torch.sort(logratios, dim=0) + ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0) + delta = logratios_sorted - ref_logratios_sorted + per_sequence_loss = ( + -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * delta) * self.label_smoothing + ) + + elif loss_type == "aot_unpaired": + chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0) + rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0) + delta = chosen_logratios_sorted - rejected_logratios_sorted + per_sequence_loss = ( + -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * delta) * self.label_smoothing + ) + + elif loss_type == "apo_zero": + # Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are better than your model's default output + # Increase chosen likelihood and decrease rejected likelihood + losses_chosen = 1 - torch.sigmoid(self.beta * chosen_logratios) + losses_rejected = torch.sigmoid(self.beta * rejected_logratios) + per_sequence_loss = losses_chosen + losses_rejected + + elif loss_type == "apo_down": + # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are worse than your model's default output. + # Decrease chosen likelihood and decrease rejected likelihood more + losses_chosen = torch.sigmoid(self.beta * chosen_logratios) + losses_rejected = 1 - torch.sigmoid(self.beta * delta_score) + per_sequence_loss = losses_chosen + losses_rejected + + elif loss_type == "discopop": + # Eqn (5) of the DiscoPOP paper (https://huggingface.co/papers/2406.08414) + logits = delta_score * self.beta + # Modulate the mixing coefficient based on the log ratio magnitudes + log_ratio_modulation = torch.sigmoid(logits / self.args.discopop_tau) + logistic_component = -F.logsigmoid(logits) + exp_component = torch.exp(-logits) + # Blend between logistic and exponential component based on log ratio modulation + per_sequence_loss = ( + logistic_component * (1 - log_ratio_modulation) + exp_component * log_ratio_modulation + ) + + elif loss_type == "sft": + chosen_logits, _ = shift_logits.chunk(2, dim=0) + chosen_labels, _ = shift_labels.chunk(2, dim=0) + chosen_mask, _ = shift_completion_mask.chunk(2, dim=0) + batch_loss = F.cross_entropy(chosen_logits[chosen_mask.bool()], chosen_labels[chosen_mask.bool()]) + # Implementation convenience: expand the scalar SFT loss to a per-sequence tensor so it matches the + # shape of other losses; only the mean is used, so this is a no-op numerically. + per_sequence_loss = batch_loss.expand(chosen_logits.size(0)) + + elif loss_type == "sigmoid_norm": + chosen_mask, rejected_mask = completion_mask.chunk(2, dim=0) + chosen_avg_score = chosen_scores / chosen_mask.sum(dim=1).clamp(min=1.0) + rejected_avg_score = rejected_scores / rejected_mask.sum(dim=1).clamp(min=1.0) + delta = chosen_avg_score - rejected_avg_score + per_sequence_loss = -F.logsigmoid(self.beta * delta) + + else: + raise ValueError( + f"Unknown loss type: {loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', " + "'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_unpaired', 'apo_zero', 'apo_down', " + "'discopop', 'sft', 'sigmoid_norm']" + ) + + if self.use_weighting: + # Eq (2) of the WPO paper: https://huggingface.co/papers/2406.11827 + completion_lengths = shift_completion_mask.sum(dim=1).clamp_min(1) + with torch.no_grad(): + lse1 = torch.logsumexp(shift_logits, dim=-1) + lse2 = torch.logsumexp(2.0 * shift_logits, dim=-1) + log_denom = lse2 - 2.0 * lse1 + aligned_logps = (per_token_logps - log_denom) * shift_completion_mask + mean_logps = aligned_logps.sum(dim=1) / completion_lengths + weights = torch.exp(mean_logps) + chosen_weights, rejected_weights = weights.chunk(2, dim=0) + per_sequence_loss *= chosen_weights * rejected_weights + + loss += per_sequence_loss.mean() * loss_weight + + # Log the metrics + # Entropy + per_token_entropy = entropy_from_logits(shift_logits.detach()) + mask = shift_completion_mask + entropy_sum = (per_token_entropy * mask).sum() + total_tokens = mask.sum() + + # Gather counts across ranks and weight-average + entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() + total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum() + entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0 + self._metrics[mode]["entropy"].append(entropy) + + # Number of tokens + if mode == "train": + num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + self._total_train_tokens += num_tokens_in_batch + self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + + # Average logits for chosen and rejected completions + chosen_logits, rejected_logits = shift_logits.detach().chunk(2, dim=0) + chosen_mask, rejected_mask = shift_completion_mask.chunk(2, dim=0) + total_chosen_logits = chosen_logits[chosen_mask.bool()].mean(-1).sum() + total_chosen_tokens = chosen_mask.sum() + total_rejected_logits = rejected_logits[rejected_mask.bool()].mean(-1).sum() + total_rejected_tokens = rejected_mask.sum() + total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item() + total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item() + total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item() + total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item() + avg_chosen_logits = total_chosen_logits / total_chosen_tokens if total_chosen_tokens > 0 else 0.0 + avg_rejected_logits = total_rejected_logits / total_rejected_tokens if total_rejected_tokens > 0 else 0.0 + self._metrics[mode]["logits/chosen"].append(avg_chosen_logits) + self._metrics[mode]["logits/rejected"].append(avg_rejected_logits) + + # Token accuracy for the chosen completions + predictions = chosen_logits.argmax(dim=-1) + chosen_mask = shift_completion_mask[: len(shift_completion_mask) // 2].bool() + chosen_labels = shift_labels[: len(shift_labels) // 2] + correct_predictions = (predictions == chosen_labels) & chosen_mask + total_tokens = chosen_mask.sum() + correct_tokens = correct_predictions.sum() + correct_tokens = self.accelerator.gather_for_metrics(correct_tokens) + total_tokens = self.accelerator.gather_for_metrics(total_tokens) + total_sum = total_tokens.sum() + accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0 + self._metrics[mode]["mean_token_accuracy"].append(accuracy) + + # Rewards for chosen and rejected completions + chosen_rewards = self.beta * chosen_logratios.detach() + rejected_rewards = self.beta * rejected_logratios.detach() + agg_chosen_rewards = self.accelerator.gather(chosen_rewards) + agg_rejected_rewards = self.accelerator.gather(rejected_rewards) + self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item()) + self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item()) + + # Reward accuracy + reward_accuracies = (chosen_rewards > rejected_rewards).float() + agg_reward_accuracies = self.accelerator.gather(reward_accuracies) + self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item()) + + # Reward margins + margins = chosen_rewards - rejected_rewards + agg_margins = self.accelerator.gather(margins) + self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item()) + + # Average log probabilities for chosen and rejected completions + self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item()) + self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item()) + + return (loss, outputs) if return_outputs else loss + + def evaluate( + self, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> dict[str, float]: + # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way + # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer. `_prepare_dataset` is + # idempotent: it skips datasets that are already tokenized. A `str` selects a dataset that was already prepared + # at init time, so it's left untouched. + if not self._is_vision_dataset and eval_dataset is not None and not isinstance(eval_dataset, str): + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, self.processing_class, self.args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, self.processing_class, self.args, "eval") + # With `precompute_ref_log_probs`, `_compute_loss` reads the reference log-probs from the batch, so they + # must be precomputed here as well, mirroring `__init__`. + if self.precompute_ref_logps: + batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size + if isinstance(eval_dataset, dict): + eval_dataset = { + name: self._precompute_ref_logps(dataset, name, batch_size) + for name, dataset in eval_dataset.items() + } + else: + eval_dataset = self._precompute_ref_logps(eval_dataset, "eval", batch_size) + return super().evaluate( + eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix + ) + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + try: + if self.use_liger_kernel: + return self._compute_loss_liger(model, inputs, return_outputs) + return self._compute_loss(model, inputs, return_outputs) + except ValueError as e: + if "Image features and image tokens do not match" in str(e) and self.args.max_length is not None: + raise ValueError( + f"The current `max_length` ({self.args.max_length}) is too short and causes image placeholder " + f"tokens in `input_ids` to be truncated, while the corresponding image features remain intact. " + f"Please increase `max_length` or set it to `None` to disable truncation." + ) from e + raise + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and + # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(), self.compute_loss_context_manager(): + if prediction_loss_only: + loss = self.compute_loss(model, inputs, return_outputs=False) # logits aren't materialized with liger + logits, labels = None, None + else: + loss, outputs = self.compute_loss(model, inputs, return_outputs=True) + logits, labels = outputs.logits, inputs["input_ids"] + return loss, logits, labels + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) +class UnslothDPOTrainer(_UnslothDPOTrainer): + """ + + Trainer for Direct Preference Optimization (DPO) method. This algorithm was initially proposed in the paper [Direct + Preference Optimization: Your Language Model is Secretly a Reward Model](https://huggingface.co/papers/2305.18290). + This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + >>> from trl import DPOTrainer + >>> from datasets import load_dataset + + >>> dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train") + + >>> trainer = DPOTrainer( + ... model="Qwen/Qwen2.5-0.5B-Instruct", + ... train_dataset=dataset, + ... ) + >>> trainer.train() + ``` + + Args: + model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `.from_pretrained` (where `` is derived from the model + config) with the keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + - A [`~peft.PeftModel`] object. Only causal language models are supported. + ref_model ([`~transformers.PreTrainedModel`], *optional*): + Reference model used to compute the reference log probabilities. + + - If provided, this model is used directly as the reference policy. + - If `None`, the trainer will automatically use the initial policy corresponding to `model`, i.e. the model + state before DPO training starts. + args ([`DPOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator ([`~transformers.DataCollator`], *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~trainer.dpo_trainer.DataCollatorForPreference`] if the model is a language model and + [`~trainer.dpo_trainer.DataCollatorForVisionPreference`] if the model is a vision-language model. Custom + collators must truncate sequences before padding; the trainer does not apply post-collation truncation. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. This trainer supports both [language modeling](#language-modeling) type and + [prompt-completion](#prompt-completion) type. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a + [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing + [`SFTConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a boolean + `compute_result` argument. This will be triggered after the last eval batch to signal that the function + needs to calculate and return the global summary statistics rather than accumulating the batch-level + statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + + """ + def __init__( + self, + model, + ref_model = None, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + compute_metrics = None, + callbacks = None, + peft_config = None, + **kwargs + ): + if args is None: args = UnslothDPOConfig() + use_bf16 = getattr(args, 'bf16', False) + if type(use_bf16) is not bool: use_bf16 = False + use_fp16 = getattr(args, 'fp16', False) + if type(use_fp16) is not bool: use_fp16 = False + force_float32 = False + full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' + if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): + print('Unsloth: Switching to float32 training since model cannot work with float16') + force_float32 = True + mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') + dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) + if dtype is None: dtype = model.get_input_embeddings().weight.dtype + from unsloth_zoo.utils import _get_dtype + dtype = _get_dtype(dtype) + float16 = dtype == torch.float16 + if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') + if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') + if force_float32: + # Forced float32 training + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': + # Mixed precision training + args.fp16 = float16 + args.bf16 = not float16 + os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' + # args.mixed_precision is a new argument which needs to be set now + elif mixed_precision_dtype == 'bfloat16': + # Both False since bfloat16 full finetuning doesn't do any autocasting. + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + + if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': + args.eval_strategy = 'steps' + if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 + ga_steps = getattr(args, 'gradient_accumulation_steps', None) + if ga_steps is not None and ga_steps > 1: + from transformers import __version__ as transformers_version + if Version(transformers_version) <= Version('4.45.2'): + print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' + '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') + if getattr(args, 'eval_strategy', 'no') != 'no': + eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) + if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size + if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps + fp16_full_eval = getattr(args, 'fp16_full_eval', False) + if type(fp16_full_eval) is not bool: fp16_full_eval = False + bf16_full_eval = getattr(args, 'bf16_full_eval', False) + if type(bf16_full_eval) is not bool: bf16_full_eval = False + if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True + if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False + if force_float32: + args.bf16_full_eval = False + args.fp16_full_eval = False + elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': + args.bf16_full_eval = True + args.fp16_full_eval = False + elif not bf16_full_eval and not fp16_full_eval: + args.bf16_full_eval = args.bf16 + args.fp16_full_eval = args.fp16 + _output_logits = False + if locals().get('compute_metrics', None) is not None: _output_logits = True + if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True + if _output_logits: + os.environ['UNSLOTH_RETURN_LOGITS'] = '1' + if model is not None: + _warnings_issued = getattr(model, 'warnings_issued', None) + if _warnings_issued is None: + model.warnings_issued = {} + elif not isinstance(_warnings_issued, dict): + try: + model.warnings_issued = dict(_warnings_issued) + except Exception: + model.warnings_issued = {} + if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): + pass + else: + model_max_seq_length = getattr(model, 'max_seq_length', None) + args_max_seq_length = getattr(args, 'max_seq_length', None) + if args_max_seq_length is None and model_max_seq_length is not None: + max_seq_length = model.max_seq_length + if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length + elif args_max_seq_length is not None and model_max_seq_length is not None: + if args_max_seq_length > model_max_seq_length: + print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' + 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') + args.max_seq_length = model_max_seq_length + if model is not None and hasattr(model, 'for_training'): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' + if 'processing_class' in locals(): + if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' + if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' + __tokenizer = processing_class if 'processing_class' in locals() else tokenizer + from unsloth_zoo.vision_utils import UnslothVisionDataCollator + if not isinstance(data_collator, UnslothVisionDataCollator): + if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names: + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names: + data_collator = DataCollatorForSeq2Seq( + __tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + else: + if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False + if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' + if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} + if not isinstance(data_collator, UnslothVisionDataCollator): + if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): + if isinstance(data_collator, DataCollatorForSeq2Seq): + data_collator = DataCollatorForSeq2Seq( + __tokenizer.tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer.tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('dpo_trainer', other_metrics) + if hasattr(train_dataset, 'column_names'): + column_names = set(train_dataset.column_names) + check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask', + 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels', + 'prompt_input_ids', 'prompt_attention_mask'] + if all(x in column_names for x in check): + train_dataset = train_dataset.remove_columns(['chosen', 'rejected', 'prompt']) + del check, column_names + if hasattr(train_dataset, 'column_names'): + column_names = set(train_dataset.column_names) + is_dpo_dataset = ({'chosen', 'rejected'}.issubset(column_names) or + {'prompt_input_ids', 'chosen_input_ids', 'rejected_input_ids'}.issubset(column_names)) + if is_dpo_dataset and isinstance(data_collator, TransformersDataCollatorForLanguageModeling): + data_collator = None + del is_dpo_dataset, column_names + from trl.trainer.dpo_trainer import DataCollatorForPreference + if not hasattr(DataCollatorForPreference, '_unsloth_vision_keys_patch'): + _old_dpo_collator_torch_call = DataCollatorForPreference.torch_call + + def _unsloth_dpo_torch_call(self, examples): + output = _old_dpo_collator_torch_call(self, examples) + import torch as _unsloth_torch + try: + from trl.trainer.utils import pad as _unsloth_trl_pad + except Exception: + _unsloth_trl_pad = None + for _k in ('pixel_position_ids', 'image_position_ids', 'mm_token_type_ids'): + if not all(_k in example for example in examples): + continue + _is_position_key = _k.endswith('position_ids') + _padding_value = -1 if _is_position_key else 0 + _padding_side = 'right' if _is_position_key else 'left' + _values = [_unsloth_torch.as_tensor(example[_k]) for example in examples] + try: + if _unsloth_trl_pad is not None: + output[_k] = _unsloth_trl_pad(_values, padding_value=_padding_value, padding_side=_padding_side) + else: + from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence + output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value) + except Exception: + from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence + output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value) + return output + + DataCollatorForPreference.torch_call = _unsloth_dpo_torch_call + DataCollatorForPreference._unsloth_vision_keys_patch = True + + # [TODO] Fix up DataParallel multiplying batch sizes + # [TODO] DDP works, but DP seems to not work? [TODO] + if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: + if getattr(args, "_n_gpu", 1) != 1: + args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + super().__init__( + model = model, + ref_model = ref_model, + args = args, + data_collator = data_collator, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + compute_metrics = compute_metrics, + callbacks = callbacks, + peft_config = peft_config,**kwargs) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() + if hasattr(self, 'neftune_hook_handle'): + self.neftune_hook_handle.remove() + if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle + if getattr(args, 'neftune_noise_alpha', None) is not None: + model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha + pass + if hasattr(self, 'accelerator'): + scaler = self.accelerator.scaler + current_model = model + while hasattr(current_model, 'model'): + current_model.accelerator_scaler = scaler + current_model = current_model.model + current_model.accelerator_scaler = scaler + pass + if hasattr(self, 'train'): + self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) + pass + if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): + _vllm_tok = self.llm.get_tokenizer() + _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) + if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: + _vllm_tok.chat_template = _pc.chat_template + pass + +pass + + +if hasattr(logger, "addFilter"): + import logging + class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not (self.text in x.getMessage()) + pass + logger.addFilter(HideLoggingMessage("`use_cache=True`")) + diff --git a/unsloth_compiled_cache/UnslothGRPOTrainer.py b/unsloth_compiled_cache/UnslothGRPOTrainer.py new file mode 100644 index 0000000000000000000000000000000000000000..c4f2191535d2f9ab86ade5b37233ecbb67ef2743 --- /dev/null +++ b/unsloth_compiled_cache/UnslothGRPOTrainer.py @@ -0,0 +1,5848 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.grpo_trainer import (Any, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, BaseTunerLayer, Callable, CommitScheduler, Dataset, DatasetCard, DatasetCardData, DistributedBackend, EnvironmentFactory, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, LoraConfig, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RepeatSampler, RewardFunc, RolloutFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, _BaseTrainer, _ForwardRedirection, add_response_schema, apply_chat_template, asyncio, atexit, copy, create_model_from_path, create_repo, defaultdict, deque, disable_dropout_in_model, disable_gradient_checkpointing, gather, gather_object, get_config_model_id, get_peft_model, get_training_chat_template, identity, inspect, is_chat_template_prefix_preserving, is_conversational, is_jmespath_available, is_liger_kernel_available, is_peft_available, is_peft_model, is_rich_available, logger, math, nanmax, nanmin, nanstd, nn, np, nullcontext, os, pad, parse_response, pd, peft, pkg_resources, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, selective_log_softmax, set_seed, shuffle_sequence_dict, shutdown_event_loop_in_daemon, split_pixel_values_by_grid, split_tensor_dict, start_event_loop_in_daemon, supports_tool_calling, sys, textwrap, time, torch, transformers, unsplit_pixel_values_by_grid, unwrap_model_for_generation, use_adapter, warnings, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, BaseTunerLayer, Callable, CommitScheduler, Dataset, DatasetCard, DatasetCardData, DistributedBackend, EnvironmentFactory, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RewardFunc, RolloutFunc, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, add_response_schema, atexit, copy, create_model_from_path, create_repo, defaultdict, deque, disable_dropout_in_model, gather, get_config_model_id, get_peft_model, get_training_chat_template, identity, inspect, is_chat_template_prefix_preserving, is_jmespath_available, is_liger_kernel_available, is_peft_available, is_peft_model, logger, nn, np, os, pad, parse_response, pd, peft, pkg_resources, prepare_deepspeed, prepare_fsdp, set_seed, shutdown_event_loop_in_daemon, start_event_loop_in_daemon, supports_tool_calling, sys, time, torch, transformers, warnings, Version, copy, gather, is_conversational, np, os, pad, parse_response, profiling_context, torch, transformers, Any, apply_chat_template, copy, disable_gradient_checkpointing, gather, gather_object, is_conversational, math, nanmax, nanmin, nanstd, np, os, pad, pd, peft, prepare_multimodal_messages, torch, use_adapter, gather, np, os, pad, profiling_context, torch, transformers, unwrap_model_for_generation, math, np, os, pad, selective_log_softmax, torch, transformers, Any, np, profiling_decorator, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch, GRPOTrainer, gather, inspect, nanmax, nanmin, np, os, pad, time, torch) + + +import os +import math +import logging +from typing import * +from dataclasses import dataclass, field +from packaging.version import Version +import torch +import numpy as np +from contextlib import nullcontext +from torch.nn import functional as F +import inspect +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling +from transformers.training_args import ParallelMode +from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize + +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass +# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain +# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. +try: + from unsloth.models._utils import _unsloth_reset_stray_compile_cache +except Exception: + def _unsloth_reset_stray_compile_cache(self): pass +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Drop any torch.compile graph cache poisoned by a stray pre-train forward. + try: + _unsloth_reset_stray_compile_cache(self) + except Exception: + pass + # Finish the previous W&B run if this is a subsequent train() call. + # We do this at the START of train() (not the end) so that + # evaluate() / log() still work after train() completes. + # HF's WandbCallback.setup() will call wandb.init() for the new run. + # See: https://github.com/unslothai/unsloth/issues/3954 + if getattr(self, '_unsloth_training_completed', False): + try: + import wandb + if wandb.run is not None: + wandb.finish() + # Reset HF's WandbCallback so it calls wandb.init() for the new run + for cb in self.callback_handler.callbacks: + if type(cb).__name__ == 'WandbCallback': + cb._initialized = False + break + except: + pass + # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + output = f(self, *args, **kwargs) + # Restore previous mode when possible + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass + # Mark that training completed so the next train() call can + # finish this W&B run before starting a new one + self._unsloth_training_completed = True + return output + return wrapper +pass + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + "triton.enable_persistent_tma_matmul": torch.cuda.get_device_capability()[0] >= 9, + "cuda.cutlass_epilogue_fusion_enabled": torch.cuda.get_device_capability()[0] >= 9, + "cuda.cutlass_tma_only": torch.cuda.get_device_capability()[0] >= 9, + "cuda.compile_opt_level" : "-O2", + "cuda.enable_cuda_lto" : True, + } + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_hidden_states_selective_log_softmax( + hidden_states: torch.Tensor, + lm_head: torch.Tensor, + index: torch.Tensor, + chunks: int = 4, + logit_scale_multiply: float = 0.0, + logit_scale_divide: float = 0.0, + logit_softcapping: float = 0.0, + temperature: float = 1.0, +) -> torch.Tensor: + # All Unsloth Zoo code licensed under AGPL3 + flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + flat_index = index.reshape(-1) + + chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) + chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) + + all_per_token_logps = [] + + for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): + chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() + + if logit_scale_multiply != 0.0: + chunk_logits = chunk_logits * logit_scale_multiply + if logit_scale_divide != 0.0: + chunk_logits = chunk_logits / logit_scale_divide + if logit_softcapping != 0.0: + chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) + + chunk_logits = chunk_logits.to(torch.float32) + + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + + selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + + all_per_token_logps = torch.concat(all_per_token_logps) + + all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) + return all_per_token_logps + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_selective_log_softmax( + logits, + index, + temperature: float = 1.0, + chunks: int = 4, +): + chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) + chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) + all_per_token_logps = [] + # Per-chunk selective_log_softmax. + for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): + chunk_logits = chunk_logits.to(torch.float32) + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + pass + all_per_token_logps = torch.concat(all_per_token_logps) + all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) + return all_per_token_logps + +def calculate_pad_tokens_in_prompt( + input_ids: torch.Tensor, + logits_to_keep: int, + pad_token_id: int +) -> torch.Tensor: + """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" + if logits_to_keep >= input_ids.shape[1]: + raise ValueError("logits_to_keep must be smaller than the sequence length.") + + prompt_section = input_ids[:, :-logits_to_keep] + + padding_mask = (prompt_section == pad_token_id) + + pad_token_counts = padding_mask.sum(dim=1) + + return pad_token_counts + +def create_completion_attention_mask( + completion_input_ids: torch.Tensor, + left_pad_tokens_per_prompt: torch.Tensor, + max_left_pad: int, + pad_token_id: int +) -> torch.Tensor: + """Build a completion mask that zeros leading prompt and trailing pad tokens. + + For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) + this returns [0,0,0,1,1,1,0,0,0]. + """ + batch_size, completion_len = completion_input_ids.shape + device = completion_input_ids.device + + num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt + + indices = torch.arange(completion_len, device=device).unsqueeze(0) + shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) + + non_padding_mask = (completion_input_ids != pad_token_id) + + final_mask = shift_mask & non_padding_mask + + return final_mask + +def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: + """Move all padding tokens in each sequence to the right.""" + mask = (tensor != pad_id) + # stable=True since the binary mask is unordered. + sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) + packed_tensor = torch.gather(tensor, 1, sorted_indices) + return packed_tensor + +def align_logprobs_with_mask( + logprob_tensor: torch.Tensor, + attention_mask: torch.Tensor, + pad_value: float = 0.0 +) -> torch.Tensor: + """Align a log probability tensor with a given attention mask.""" + + device = logprob_tensor.device + batch_size, logprob_seq_len = logprob_tensor.shape + mask_seq_len = attention_mask.shape[1] + + padded_logprobs = torch.full( + attention_mask.shape, + fill_value=pad_value, + dtype=logprob_tensor.dtype, + device=device + ) + + left_pad_counts = torch.argmax(attention_mask, dim=1) + + cols = torch.arange(logprob_seq_len, device=device) + dest_indices = left_pad_counts.unsqueeze(1) + cols + + # Destination row indices, shape [batch_size, logprob_seq_len]. + row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) + + # Keep only in-bounds destinations, then scatter via advanced indexing. + valid_mask = dest_indices < mask_seq_len + valid_rows = row_indices[valid_mask] + valid_cols = dest_indices[valid_mask] + valid_vals = logprob_tensor[valid_mask] + padded_logprobs[valid_rows, valid_cols] = valid_vals + + return padded_logprobs + +def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" + if tool_mask is None: + return completion_mask + if tool_mask.shape[0] != completion_mask.shape[0]: + raise ValueError("tool_mask batch size must match completion_mask batch size.") + + tool_mask = tool_mask.to(device=completion_mask.device) + if tool_mask.shape == completion_mask.shape: + aligned_tool_mask = tool_mask + else: + aligned_tool_mask = align_logprobs_with_mask( + tool_mask, + completion_mask, + pad_value=0, + ) + return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) + +def autotune_batch_and_chunks( + total_input_rows, + seq_len, + hidden_size, + vocab_size, + dtype_bytes=16, + multiplier=None +): + if multiplier is None: + final_m = max(4, seq_len // 4096) + else: + final_m = multiplier + + if torch.cuda.is_available(): + free_bytes, _ = torch.cuda.mem_get_info() + limit_gb = (free_bytes / (1024**3))*.80 + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + # XPU: estimate free memory as total - reserved. + total_mem = torch.xpu.get_device_properties(0).total_memory + reserved_mem = torch.xpu.memory_reserved() + free_bytes = total_mem - reserved_mem + limit_gb = (free_bytes / (1024**3)) * 0.80 + else: + # Fallback: assume 8GB available. + limit_gb = 8.0 + + bytes_to_gb = 1024**3 + + b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) + + hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb + + base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb + logits_gb = base_logits / final_m + + total_mem_gb = hidden_gb + logits_gb + + valid_mask = total_mem_gb <= limit_gb + valid_indices = torch.nonzero(valid_mask, as_tuple=False) + + if valid_indices.shape[0] == 0: + #This means your GPU will OOM + return 4, final_m + + best_idx = valid_indices[0].item() + final_b = int(b_vals[best_idx].item()) + + return final_b, final_m + +def sanitize_logprob(logprob): + """Local port of trl.scripts.vllm_serve.sanitize_logprob. + Filters NaN logprobs from vLLM outputs.""" + value = logprob.logprob + if math.isnan(value): + logging.getLogger(__name__).warning( + f"Generated NaN logprob, token logprob '{logprob}' will be ignored" + ) + return None + return value +def _unsloth_get_final_logit_softcapping(config): + """Return final_logit_softcapping for a model config, falling back to the + nested text sub-config for composite models. Handles both: + - Gemma-4-style configs where the attribute lives on ``config.text_config`` + - T5Gemma-style composite configs where the text sub-config is only + reachable via ``config.get_text_config()`` + Returns 0 if unset, matching the previous behaviour. + """ + softcap = getattr(config, "final_logit_softcapping", None) + if softcap is None: + text_cfg = getattr(config, "text_config", None) + if text_cfg is None: + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_cfg = get_text_config() + except (TypeError, ValueError): + text_cfg = None + if text_cfg is not None and text_cfg is not config: + softcap = getattr(text_cfg, "final_logit_softcapping", None) + return 0 if softcap is None else softcap + +def _unsloth_get_mm_token_id(processing_class, attr_name, token): + tokenizer = getattr(processing_class, "tokenizer", processing_class) + token_id = getattr(processing_class, attr_name, None) + if token_id is None: + token_id = getattr(tokenizer, attr_name, None) + + convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None) + if token_id is None and convert_tokens_to_ids is not None: + token_id = convert_tokens_to_ids(token) + + if type(token_id) is int and token_id >= 0: + if token_id != getattr(tokenizer, "unk_token_id", None): + return token_id + return None + +def _unsloth_fix_mm_token_type_ids( + processing_class, input_ids, mm_token_type_ids = None, completion_ids = None +): + image_token_id = _unsloth_get_mm_token_id( + processing_class, "image_token_id", "<|image_pad|>" + ) + video_token_id = _unsloth_get_mm_token_id( + processing_class, "video_token_id", "<|video_pad|>" + ) + + if image_token_id is not None or video_token_id is not None: + rebuilt = input_ids.new_zeros(input_ids.shape) + if image_token_id is not None: + rebuilt = rebuilt.masked_fill(input_ids == image_token_id, 1) + if video_token_id is not None: + rebuilt = rebuilt.masked_fill(input_ids == video_token_id, 2) + return rebuilt + + if ( + mm_token_type_ids is not None + and completion_ids is not None + and mm_token_type_ids.shape[0] == input_ids.shape[0] + and mm_token_type_ids.shape[1] + completion_ids.shape[1] == input_ids.shape[1] + ): + return torch.cat( + [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], + dim = 1, + ) + return mm_token_type_ids + +def _unsloth_clear_stateful_mrope(model): + modules = getattr(model, "modules", None) + if modules is None: + return False + + cleared = False + for module in modules(): + if hasattr(module, "compute_3d_position_ids") and hasattr(module, "rope_deltas"): + module.rope_deltas = None + cleared = True + return cleared + +def grpo_compute_loss( + ref, + new, + old, + sampling_per_token_logps, + input_ids, + mask, + beta, + advantages, + **kwargs +): + # All Unsloth Zoo code licensed under AGPL3 + # Optional argument defaults. + loss_type = kwargs.get("loss_type", "grpo") + epsilon_low = kwargs.get("epsilon_low", 0.2) + epsilon_high = kwargs.get("epsilon_high", 0.2) + max_completion_length = kwargs.get("max_completion_length", 8192) + delta = kwargs.get("delta", None) + importance_sampling_level = kwargs.get("importance_sampling_level", "token") + num_items_in_batch = kwargs.get("num_items_in_batch", None) + current_gradient_accumulation_steps = kwargs.get("current_gradient_accumulation_steps", 1) + num_processes = kwargs.get("num_processes", 1) + use_vllm = kwargs.get("use_vllm", False) + vllm_importance_sampling_cap = kwargs.get("vllm_importance_sampling_cap", 2.0) + get_sapo_token_loss = kwargs.get("get_sapo_token_loss", None) + sapo_temperature_pos = kwargs.get("sapo_temperature_pos", 1.0) + sapo_temperature_neg = kwargs.get("sapo_temperature_neg", 1.05) + get_gamma_weights = kwargs.get("get_gamma_weights", None) + vespo_k_pos = kwargs.get("vespo_k_pos", 2.0) + vespo_lambda_pos = kwargs.get("vespo_lambda_pos", 3.0) + vespo_k_neg = kwargs.get("vespo_k_neg", 3.0) + vespo_lambda_neg = kwargs.get("vespo_lambda_neg", 2.0) + get_off_policy_mask = kwargs.get("get_off_policy_mask", None) + off_policy_mask_threshold = kwargs.get("off_policy_mask_threshold", None) + input_ids = input_ids.unsqueeze(-1) + + if advantages.dim() == 1: + advantages = advantages.unsqueeze(1) + + if off_policy_mask_threshold is not None: + off_policy_mask = get_off_policy_mask( + advantages=advantages, + per_token_logps=new, + old_per_token_logps=old, + mask=mask, + off_policy_threshold=off_policy_mask_threshold, + ) + + with torch.no_grad(): + if use_vllm and sampling_per_token_logps is not None: + # Filter out extra leading prompt tokens after left-padding input_ids. + importance_sampling_ratio = torch.exp((old * mask) - sampling_per_token_logps) + importance_sampling_ratio = torch.clamp( + importance_sampling_ratio, max=vllm_importance_sampling_cap + ) + pass + + # Must detach when old is None: exp(new - new.detach()) == 1 but keeps grads correct. + if old is not None: + log_ratio = new - old + else: + log_ratio = new - new.detach() + + if importance_sampling_level == "token": + log_importance_weights = log_ratio + elif importance_sampling_level == "sequence": + log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0) + log_importance_weights = log_importance_weights.unsqueeze(-1) + else: + raise ValueError( + f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' " + "and 'sequence'." + ) + + coef_1 = torch.exp(log_importance_weights) + + # Reverse KL: low-variance low-bias estimator as used in the GRPO paper. + if beta != 0.0: + kl_i = torch.exp(ref - new) - (ref - new) - 1.0 + + else: + # Zeros with the correct shape. + if importance_sampling_level == "sequence": + kl_i = new.new_zeros(new.size(0), 1) + else: + kl_i = torch.zeros_like(new) + + if loss_type == "cispo": + clamped_ratios = torch.clamp(coef_1, max=epsilon_high).detach() + loss_i = -clamped_ratios * advantages * new + elif loss_type in ["grpo", "bnpo", "dr_grpo", "dapo"]: + coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high) + + if delta is not None: + loss_1 = torch.clamp(coef_1, max=delta) * advantages + else: + loss_1 = coef_1 * advantages + pass + loss_2 = coef_2 * advantages + loss_i = -torch.min(loss_1, loss_2) + elif loss_type == "sapo": + if get_sapo_token_loss is None: + raise Exception(f"sapo is only available in TRL 0.26.0+") + loss_i = torch.empty_like(coef_1) + positive_advantages_mask = advantages.repeat([1, coef_1.shape[1]]) > 0 + # With n_chunks some tensors may be empty; guard the indexing. + if coef_1[positive_advantages_mask].numel() != 0: + loss_i[positive_advantages_mask] = get_sapo_token_loss( + coef_1[positive_advantages_mask], sapo_temperature_pos + ) + if coef_1[~positive_advantages_mask].numel() != 0: + loss_i[~positive_advantages_mask] = get_sapo_token_loss( + coef_1[~positive_advantages_mask], sapo_temperature_neg + ) + loss_i = -loss_i * advantages + elif loss_type == "vespo": + if get_gamma_weights is None: + raise Exception("vespo is only available in TRL 0.26.0+") + phi_seq = get_gamma_weights( + advantages=advantages, + log_ratio_per_token=log_ratio, + mask=mask, + importance_sampling_ratio=kwargs.get("importance_sampling_ratio"), + k_pos=vespo_k_pos, + lambda_pos=vespo_lambda_pos, + k_neg=vespo_k_neg, + lambda_neg=vespo_lambda_neg, + ) + loss_i = -phi_seq * advantages * new + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + if off_policy_mask_threshold is not None: + loss_i = loss_i * off_policy_mask + + if use_vllm and sampling_per_token_logps is not None: + loss_i = loss_i * importance_sampling_ratio + # delta for the metric. + with torch.no_grad(): + delta = torch.abs(old - sampling_per_token_logps) + delta = delta * mask + flat_is_ratio = importance_sampling_ratio * mask + else: + delta = torch.tensor([]).detach() + flat_is_ratio = torch.tensor([]).detach() + if beta != 0.0: + loss_i = loss_i + beta * kl_i + + mask = mask.to(torch.float32) + n_mask_per_reward = mask.sum(1) + + # https://github.com/huggingface/trl/blob/e8b8499f1f8d76838155b515e414ee98f757d6d5/trl/trainer/grpo_trainer.py#L1624 + if loss_type in ["grpo", "sapo"]: + loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() + loss = loss / current_gradient_accumulation_steps + elif loss_type == "bnpo": + loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0) + loss = loss / current_gradient_accumulation_steps + elif loss_type == "dr_grpo": + loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length) + loss = loss / current_gradient_accumulation_steps + elif loss_type in ["cispo", "dapo", "vespo"]: + normalizer = num_items_in_batch/ num_processes + loss = (loss_i * mask).sum() / normalizer + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + # Folded metrics. + def masked_batch_mean(x): + with torch.inference_mode(): + completion_length = n_mask_per_reward.mean() + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return completion_length, x.mean() + else: + mean_kl_per_reward = (x * mask).sum(1) / n_mask_per_reward + mean_kl = mean_kl_per_reward.mean() + return completion_length, mean_kl + completion_length, mean_kl = masked_batch_mean(kl_i) + return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, mask + +class UnslothEfficientGRPO(torch.autograd.Function): + # All Unsloth Zoo code licensed under AGPL3 + @staticmethod + def forward(ctx, _new_logps, _old_logps, _ref_logps, _sampling_per_token_logps, lm_head, _input_ids, _mask, _advantages, beta, scaler = None, n_chunks = 1, extra_kwargs=None): + if extra_kwargs is None: + extra_kwargs = {} + def compute_loss(new_logps, old_logps, ref_logps, sampling_per_token_logps, input_ids, mask, advantages, scaling): + loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, _mask = grpo_compute_loss( + ref_logps, + new_logps, + old_logps, + sampling_per_token_logps, + input_ids, + mask, + beta, + advantages, + **extra_kwargs, + ) + + # Scale for mixed precision; return loss.detach() or autograd uses 2x VRAM. + scaled_loss = loss * scaling + return scaled_loss, (loss.detach(), completion_length, mean_kl, delta, flat_is_ratio, coef_1) + pass + + device =_new_logps.device + grad_inputs = torch.empty_like(_new_logps) + accumulated_loss = torch.zeros(1, device = device)[0] + accumulated_completion_length = torch.zeros(1, device = device)[0] + accumulated_mean_kl = torch.zeros(1, device = device)[0] + accumulated_delta = [] + accumulated_flat_is_ratio = [] + accumulated_coef_1 = [] + + def accumulate_chunk( + new_logps_j, + old_logps_j, + ref_logps_j, + sampling_per_token_logps_j, + input_ids_j, + mask_j, + advantages_j, + scaling, + grad_inputs_j, + ): + (chunk_grad_input,), (chunk_loss, (unscaled_loss, chunk_completion_length, chunk_mean_kl, chunk_delta, chunk_flat_is_ratio, chunk_coef_1)) = torch.func.grad_and_value( + compute_loss, + argnums = (0,), + has_aux = True, + )(new_logps_j, old_logps_j, ref_logps_j, sampling_per_token_logps_j, input_ids_j, mask_j, advantages_j, scaling) + accumulated_loss .add_(unscaled_loss) + accumulated_completion_length.add_(chunk_completion_length) + accumulated_mean_kl .add_(chunk_mean_kl) + accumulated_delta .append(chunk_delta) + accumulated_flat_is_ratio .append(chunk_flat_is_ratio) + accumulated_coef_1 .append(chunk_coef_1) + grad_inputs_j[:] = chunk_grad_input + pass + + accumulate_chunk = torch.compile( + accumulate_chunk, + fullgraph = True, + # [TODO] Dynamic marking causes torch.compile errors if sequence length is long + dynamic = True, + options = torch_compile_options, + ) + + grad_inputs_chunks = torch.chunk(grad_inputs, chunks = n_chunks, dim = 0) + new_logps = torch.chunk(_new_logps, chunks = n_chunks, dim = 0) + if _old_logps is not None: + old_logps = torch.chunk(_old_logps, chunks = n_chunks, dim = 0) + else: + old_logps = [None] * n_chunks + if _ref_logps is not None: + ref_logps = torch.chunk(_ref_logps, chunks = n_chunks, dim = 0) + else: + ref_logps = [None] * n_chunks + if _sampling_per_token_logps is not None: + sampling_per_token_logps = torch.chunk(_sampling_per_token_logps, chunks = n_chunks, dim = 0) + else: + sampling_per_token_logps = [None] * n_chunks + input_ids = torch.chunk(_input_ids, chunks = n_chunks, dim = 0) + mask = torch.chunk(_mask, chunks = n_chunks, dim = 0) + advantages = torch.chunk(_advantages, chunks = n_chunks, dim = 0) + + # Mixed precision scaling if present. + scaling = scaler.get_scale() if scaler is not None else 1.0 + + for (grad_inputs_j, new_logps_j, old_logps_j, ref_logps_j, sampling_per_token_logps_j, input_ids_j, mask_j, advantages_j, ) in \ + zip(grad_inputs_chunks, new_logps, old_logps, ref_logps, sampling_per_token_logps, input_ids, mask, advantages): + + # [TODO] Dynamic marking causes torch.compile errors if sequence length is long + + # mark_dynamic(new_hidden_states_j) + # mark_dynamic(ref_hidden_states_j) + # if old_hidden_states_j is not None: + # mark_dynamic(old_hidden_states_j) + # mark_dynamic(input_ids_j) + # mark_dynamic(mask_j) + accumulate_chunk( + new_logps_j, + old_logps_j, + ref_logps_j, + sampling_per_token_logps_j, + input_ids_j, + mask_j, + advantages_j, + scaling, + grad_inputs_j, + ) + pass + + grad_inputs .div_(n_chunks) + accumulated_loss .div_(n_chunks) + accumulated_completion_length.div_(n_chunks) + accumulated_mean_kl .div_(n_chunks) + + if _sampling_per_token_logps is not None: + accumulated_delta = torch.cat(accumulated_delta, dim=0) + accumulated_flat_is_ratio = torch.cat(accumulated_flat_is_ratio, dim=0) + else: + accumulated_delta = None + accumulated_flat_is_ratio = None + accumulated_coef_1 = torch.cat(accumulated_coef_1, dim=0) + ctx.save_for_backward(grad_inputs) + return ( + accumulated_loss, + accumulated_completion_length, + accumulated_mean_kl, + accumulated_delta, + accumulated_flat_is_ratio, + accumulated_coef_1 + ) + pass + + @staticmethod + def backward(ctx, grad_output, dcompletion_length, dmean_kl, ddelta, ddflat_is_ratio, dcoef_1): + (grad_input,) = ctx.saved_tensors + return (grad_input, None, None, None, None, None, None, None, None, None, None, None) + pass + +def grpo_accumulated_loss( + trainer, + input_ids, + attention_mask, + logits_to_keep, + completion_mask, + advantages, + old_logps, + ref_logps, + n_chunks = -1, + tool_mask = None, + **kwargs, +): + # All Unsloth Zoo code licensed under AGPL3 + bsz, qlen = input_ids.shape + + pixel_values = kwargs.get('pixel_values',None) + image_grid_thw = kwargs.get('image_grid_thw',None) + pixel_attention_mask = kwargs.get('pixel_attention_mask',None) + image_sizes = kwargs.get('image_sizes',None) + num_images = kwargs.get('num_images',None) + # Transformers 5.x requires token_type_ids/mm_token_type_ids for some vision models + token_type_ids = kwargs.get('token_type_ids',None) + mm_token_type_ids = kwargs.get('mm_token_type_ids',None) + if mm_token_type_ids is not None or image_grid_thw is not None: + mm_token_type_ids = _unsloth_fix_mm_token_type_ids( + trainer.processing_class, input_ids, mm_token_type_ids + ) + sampling_per_token_logps = kwargs.get("sampling_per_token_logps", None) if getattr(trainer, "vllm_importance_sampling_correction", False) else None + temperature = kwargs.get("temperature", 1.0) + logit_scale_multiply = kwargs.get("logit_scale_multiply", 0.0) + logit_scale_divide = kwargs.get("logit_scale_divide", 0.0) + logit_softcapping = kwargs.get("logit_softcapping", 0.0) + prev_max_left_pad = kwargs.get("max_left_pad", 0) # max_left_pad for LLM training, enabled by default. + + # Pop from kwargs to avoid downstream issues. + _ = kwargs.pop("sampling_per_token_logps", None) + kwargs["vllm_importance_sampling_cap"] = trainer.vllm_importance_sampling_cap if sampling_per_token_logps is not None else None + kwargs["get_sapo_token_loss"] = trainer.get_sapo_token_loss if hasattr(trainer, "get_sapo_token_loss") else None + kwargs["sapo_temperature_pos"] = trainer.args.sapo_temperature_pos if hasattr(trainer.args, "sapo_temperature_pos") else None + kwargs["sapo_temperature_neg"] = trainer.args.sapo_temperature_neg if hasattr(trainer.args, "sapo_temperature_neg") else None + kwargs["get_gamma_weights"] = trainer.get_gamma_weights if hasattr(trainer, "get_gamma_weights") else None + kwargs["vespo_k_pos"] = trainer.args.vespo_k_pos if hasattr(trainer.args, "vespo_k_pos") else 2.0 + kwargs["vespo_k_neg"] = trainer.args.vespo_k_neg if hasattr(trainer.args, "vespo_k_neg") else 3.0 + kwargs["vespo_lambda_pos"] = trainer.args.vespo_lambda_pos if hasattr(trainer.args, "vespo_lambda_pos") else 3.0 + kwargs["vespo_lambda_neg"] = trainer.args.vespo_lambda_neg if hasattr(trainer.args, "vespo_lambda_neg") else 2.0 + kwargs["get_off_policy_mask"] = trainer.get_off_policy_mask if hasattr(trainer, "get_off_policy_mask") else None + kwargs["off_policy_mask_threshold"] = trainer.args.off_policy_mask_threshold if hasattr(trainer.args, "off_policy_mask_threshold") else None + kwargs["use_vllm"] = trainer.use_vllm + # Snap n_chunks to the closest divisor of bsz. + factors = [i for i in range(1, bsz + 1) if bsz % i == 0] + if n_chunks == -1: n_chunks = bsz + n_chunks = factors[min(np.searchsorted(factors, n_chunks), len(factors)-1)] + + if not hasattr(trainer, '_autocast_dtype'): + trainer._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16 + if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': trainer._autocast_dtype = None + pass + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + + lm_head = trainer.model.get_output_embeddings().weight + dtype_bytes = 16 if trainer._autocast_dtype in [torch.float16, torch.bfloat16] else 32 + + total_rows = input_ids.shape[0] + seq_len = input_ids.shape[1] + hidden_dim = lm_head.shape[1] + vocab_dim = lm_head.shape[0] + + if trainer.args.unsloth_grpo_mini_batch is None: + if not hasattr(trainer, "_has_autotuned"): + trainer._has_autotuned = True + B, multiplier = autotune_batch_and_chunks( + total_rows, seq_len, hidden_dim, vocab_dim, dtype_bytes, trainer.args.unsloth_logit_chunk_multiplier + ) + trainer.args.unsloth_grpo_mini_batch = max(1, total_rows//B) + trainer.args.unsloth_logit_chunk_multiplier = multiplier + B = trainer.args.unsloth_grpo_mini_batch + multiplier = trainer.args.unsloth_logit_chunk_multiplier + elif trainer._step % trainer.current_gradient_accumulation_steps == 0: + B = trainer.args.unsloth_grpo_mini_batch + multiplier = trainer.args.unsloth_logit_chunk_multiplier + del trainer._has_autotuned + del trainer.args.unsloth_grpo_mini_batch + del trainer.args.unsloth_logit_chunk_multiplier + else: + B = trainer.unsloth_grpo_mini_batch + multiplier = trainer.args.unsloth_logit_chunk_multiplier + else: + if trainer.args.unsloth_grpo_mini_batch > total_rows: + B = total_rows + else: + B = trainer.args.unsloth_grpo_mini_batch + + if trainer.args.unsloth_logit_chunk_multiplier is None: + multiplier = max(4, seq_len // 4096) + else: + multiplier = trainer.args.unsloth_logit_chunk_multiplier + + if pixel_values is None: + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(input_ids, logits_to_keep, trainer.processing_class.pad_token_id) + + # Determine max_left_pad from precomputed logprobs shape for consistency + if old_logps is not None: + max_left_pad = old_logps.shape[1] - logits_to_keep + elif ref_logps is not None: + max_left_pad = ref_logps.shape[1] - logits_to_keep + else: + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() + + input_ids = left_pack_padding(input_ids, trainer.processing_class.pad_token_id) + + completion_input_ids = input_ids[:, -(logits_to_keep +max_left_pad):] + completion_mask = create_completion_attention_mask(completion_input_ids, left_pad_tokens_per_prompt, max_left_pad, trainer.processing_class.pad_token_id).to(attention_mask.dtype) + + if trainer.use_vllm and sampling_per_token_logps is not None and getattr(trainer, "vllm_importance_sampling_correction", False): + sampling_per_token_logps = align_logprobs_with_mask(sampling_per_token_logps, completion_mask) + else: + sampling_per_token_logps = None + completion_mask = align_completion_tool_mask(tool_mask, completion_mask) + attention_mask = input_ids != trainer.processing_class.pad_token_id + attention_mask = attention_mask.to(attention_mask.dtype) + else: + completion_input_ids = input_ids[:, -logits_to_keep:] + completion_mask = align_completion_tool_mask(tool_mask, completion_mask) + + unwrapped_model = trainer.accelerator.unwrap_model(trainer.model, keep_fp32_wrapper = False) + + for module in unwrapped_model.modules(): + if hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "io_same_decice"): + module._hf_hook.io_same_decice = False + pass + + all_logprobs_list = [] + + def slice_sample_axis(value, start, end): + if value is None: + return None + return value[start:end] + + import math + total_samples = input_ids.shape[0] + batch_size = math.ceil(total_samples / B) + if isinstance(num_images, torch.Tensor): + num_images = num_images.detach().cpu().reshape(-1).tolist() + if image_grid_thw is not None and pixel_values is not None and num_images is not None: + rows_per_image = image_grid_thw.prod(dim=-1) + rows_per_sample = torch.split(rows_per_image, num_images) + rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) + cum_rows = torch.cat( + [ + torch.tensor([0], device=rows_per_sample.device), + rows_per_sample.cumsum(0), + ] + ) + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + else: + cum_rows = None + cum_imgs = None + + input_ids_chunks = [] + attention_mask_chunks = [] + completion_ids_chunks = [] + pixel_values_chunks = [] + image_grid_thw_chunks = [] + pixel_attention_mask_chunks = [] + image_sizes_chunks = [] + token_type_ids_chunks = [] + mm_token_type_ids_chunks = [] + + current_pixel_idx = 0 + #TRL 0.23.0 batching logic + for start in range(0, total_samples, batch_size): + end = min(start + batch_size, total_samples) + + input_ids_chunks.append(input_ids[start:end]) + attention_mask_chunks.append(attention_mask[start:end]) + completion_ids_chunks.append(completion_input_ids[start:end]) + image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) + token_type_ids_chunks.append(slice_sample_axis(token_type_ids, start, end)) + mm_token_type_ids_chunks.append( + slice_sample_axis(mm_token_type_ids, start, end) + ) + + if image_grid_thw is not None and pixel_values is not None: + + if num_images is None: + grid_slice = image_grid_thw[start:end] + batch_pixel_count = grid_slice.prod(dim=-1).sum().item() + start_pixel_idx = current_pixel_idx + end_pixel_idx = current_pixel_idx + batch_pixel_count + current_pixel_idx = end_pixel_idx + else: + start_pixel_idx = cum_rows[start].item() + end_pixel_idx = cum_rows[end].item() + img_start, img_end = cum_imgs[start], cum_imgs[end] + grid_slice = image_grid_thw[img_start:img_end] + image_grid_thw_chunks.append(grid_slice) + + pixel_values_chunks.append(pixel_values[start_pixel_idx:end_pixel_idx]) + + if pixel_attention_mask is not None: + if pixel_attention_mask.shape[0] == pixel_values.shape[0]: + pixel_attention_mask_chunks.append(pixel_attention_mask[start_pixel_idx:end_pixel_idx]) + else: + pixel_attention_mask_chunks.append(pixel_attention_mask[start:end]) + else: + pixel_attention_mask_chunks.append(None) + + else: + pixel_values_chunks.append(None) + image_grid_thw_chunks.append(None) + pixel_attention_mask_chunks.append(None) + + zipped_inputs = zip( + input_ids_chunks, + attention_mask_chunks, + pixel_values_chunks, + image_grid_thw_chunks, + pixel_attention_mask_chunks, + image_sizes_chunks, + token_type_ids_chunks, + mm_token_type_ids_chunks, + completion_ids_chunks + ) + + if trainer._autocast_dtype is None: + autocaster = nullcontext() + else: + autocaster = torch.amp.autocast(device_type = trainer.model.device.type, dtype = trainer._autocast_dtype) + + def to_device(tensor, device, non_blocking=True): + if tensor is None: return None + return tensor.to(device, non_blocking=non_blocking) + + class Unsloth_Offloaded_Log_Softmax(torch.autograd.Function): + """Manual gradient checkpointing / CPU offloading for log softmax.""" + @staticmethod + def forward(ctx, hidden_states, lm_head, index, chunks, + logit_scale_multiply, logit_scale_divide, + logit_softcapping, temperature): + # Detach so we don't keep the graph (and extra memory) on CPU. + ctx.saved_hidden_states = hidden_states.detach().contiguous().to("cpu", non_blocking=True) + ctx.device = hidden_states.device + ctx.dtype = hidden_states.dtype + + ctx.lm_head = lm_head + ctx.lm_head_requires_grad = lm_head.requires_grad + ctx.index = index + ctx.args = (chunks, logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature) + + with torch.no_grad(): + output = chunked_hidden_states_selective_log_softmax( + hidden_states, lm_head, index, *ctx.args + ) + + return output + + @staticmethod + def backward(ctx, grad_output): + hidden_states = to_device(ctx.saved_hidden_states, ctx.device) + hidden_states = hidden_states.to(ctx.dtype) + hidden_states.requires_grad_(True) + + lm_head = ctx.lm_head + # #Possibly redundant lines + # if ctx.lm_head_requires_grad: + # hidden_states.requires_grad_(True) + # else: + # lm_head = lm_head.detach() + + index = ctx.index + + with torch.enable_grad(): + output = chunked_hidden_states_selective_log_softmax( + hidden_states, lm_head, index, *ctx.args + ) + + torch.autograd.backward(output, grad_output) + + return ( + hidden_states.grad, + lm_head.grad if ctx.lm_head_requires_grad else None, + None, + None, + None, + None, + None, + None, + ) + + def efficient_log_softmax(hidden_states, lm_head, index, chunks=32, + logit_scale_multiply=0.0, logit_scale_divide=0.0, + logit_softcapping=0.0, temperature=1, batch_size=8): + if (index.shape[1] <= 1024 and batch_size <= 8) or batch_size==1: + # Normal path is faster / saves a GB under these conditions. + return chunked_hidden_states_selective_log_softmax( + hidden_states, + lm_head, + index, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature + ) + else: + return Unsloth_Offloaded_Log_Softmax.apply( + hidden_states, lm_head, index, chunks, + logit_scale_multiply, logit_scale_divide, + logit_softcapping, temperature + ) + + def compute_logprobs_chunk(new_hidden_states_chunk, completion_ids, input_ids_chunk): + # Hidden states -> lm_head matmul path; raw logits -> skip matmul and + # skip scale/softcap (model forward already applied them). + chunks = input_ids_chunk.shape[0] * multiplier + if new_hidden_states_chunk.shape[-1] == lm_head.shape[1]: + return efficient_log_softmax( + new_hidden_states_chunk, + lm_head, + completion_ids, + chunks = chunks, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + logit_softcapping = logit_softcapping, + temperature = temperature, + batch_size = B, + ) + return chunked_selective_log_softmax( + new_hidden_states_chunk, + completion_ids, + temperature = temperature, + chunks = chunks, + ) + for ( + input_ids_chunk, + attention_mask_chunk, + pixel_values_chunk, + image_grid_thw_chunk, + pixel_attention_mask_chunk, + image_sizes_chunk, + token_type_ids_chunk, + mm_token_type_ids_chunk, + completion_ids + ) in zipped_inputs: + _extra_vision_kwargs = {} + if token_type_ids_chunk is not None: + _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk + if mm_token_type_ids_chunk is not None: + _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk + with autocaster: + if pixel_values is None: + new_hidden_states_chunk = unwrapped_model( + input_ids = input_ids_chunk, + attention_mask = attention_mask_chunk, + pixel_values = pixel_values_chunk, + image_grid_thw = image_grid_thw_chunk, + pixel_attention_mask = pixel_attention_mask_chunk, + image_sizes = image_sizes_chunk, + **_extra_vision_kwargs, + ).logits + + new_hidden_states_chunk = new_hidden_states_chunk[:, -(logits_to_keep + max_left_pad + 1): , :] + new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :] + logprobs_chunk = compute_logprobs_chunk(new_hidden_states_chunk, completion_ids, input_ids_chunk) + else: + new_hidden_states_chunk = unwrapped_model( + input_ids = input_ids_chunk, + attention_mask = attention_mask_chunk, + pixel_values = pixel_values_chunk, + image_grid_thw = image_grid_thw_chunk, + pixel_attention_mask = pixel_attention_mask_chunk, + image_sizes = image_sizes_chunk, + logits_to_keep = logits_to_keep + 1, + **_extra_vision_kwargs, + ).logits + + new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :] + logprobs_chunk = compute_logprobs_chunk(new_hidden_states_chunk, completion_ids, input_ids_chunk) + # Avoids race conditions with GPT OSS offload_embbed=True; no measurable slowdown. + device_synchronize() + all_logprobs_list.append(logprobs_chunk) + + new_logprobs = torch.cat(all_logprobs_list, dim=0) + + with autocaster: + loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1 = UnslothEfficientGRPO.apply( + new_logprobs, + old_logps, + ref_logps, + sampling_per_token_logps, + lm_head, + completion_input_ids, + completion_mask, + advantages, + trainer.beta, + trainer.accelerator.scaler, + 1, + kwargs + ) + + # Force logits (not hidden states) again or output is gibberish. + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" + + return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, completion_mask + # Old non-efficient code path (dead). + new_logits = torch.matmul(new_hidden_states, lm_head.t()) + new_logits = new_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + old_logits = torch.matmul(old_hidden_states, lm_head.t()) + old_logits = old_logits[:, :-1, :] # exclude the last logit: it corresponds to the next token pred + loss, completion_length, mean_kl = grpo_compute_loss( + old_logits, + new_logits, + completion_input_ids, + completion_mask, + trainer.beta, + advantages, + ) + return loss, completion_length, mean_kl + pass + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options) +def grpo_compute_loss_slow( + ref, + new, + old, + sampling_per_token_logps, + input_ids, + mask, + beta, + advantages, + **kwargs +): + # All Unsloth Zoo code licensed under AGPL3 + # Optional argument defaults. + loss_type = kwargs.get("loss_type", "grpo") + epsilon_low = kwargs.get("epsilon_low", 0.2) + epsilon_high = kwargs.get("epsilon_high", 0.2) + max_completion_length = kwargs.get("max_completion_length", 8192) + delta = kwargs.get("delta", None) + importance_sampling_level = kwargs.get("importance_sampling_level", "token") + num_items_in_batch = kwargs.get("num_items_in_batch", None) + current_gradient_accumulation_steps = kwargs.get("current_gradient_accumulation_steps", 1) + num_processes = kwargs.get("num_processes", 1) + use_vllm = kwargs.get("use_vllm", False) + vllm_importance_sampling_cap = kwargs.get("vllm_importance_sampling_cap", 2.0) + get_sapo_token_loss = kwargs.get("get_sapo_token_loss", None) + sapo_temperature_pos = kwargs.get("sapo_temperature_pos", 1.0) + sapo_temperature_neg = kwargs.get("sapo_temperature_neg", 1.05) + get_gamma_weights = kwargs.get("get_gamma_weights", None) + vespo_k_pos = kwargs.get("vespo_k_pos", 2.0) + vespo_lambda_pos = kwargs.get("vespo_lambda_pos", 3.0) + vespo_k_neg = kwargs.get("vespo_k_neg", 3.0) + vespo_lambda_neg = kwargs.get("vespo_lambda_neg", 2.0) + get_off_policy_mask = kwargs.get("get_off_policy_mask", None) + off_policy_mask_threshold = kwargs.get("off_policy_mask_threshold", None) + input_ids = input_ids.unsqueeze(-1) + + if advantages.dim() == 1: + advantages = advantages.unsqueeze(1) + + if off_policy_mask_threshold is not None: + off_policy_mask = get_off_policy_mask( + advantages=advantages, + per_token_logps=new, + old_per_token_logps=old, + mask=mask, + off_policy_threshold=off_policy_mask_threshold, + ) + + with torch.no_grad(): + if use_vllm and sampling_per_token_logps is not None: + # Filter out extra leading prompt tokens after left-padding input_ids. + importance_sampling_ratio = torch.exp((old * mask) - sampling_per_token_logps) + importance_sampling_ratio = torch.clamp( + importance_sampling_ratio, max=vllm_importance_sampling_cap + ) + pass + + # Must detach when old is None: exp(new - new.detach()) == 1 but keeps grads correct. + if old is not None: + log_ratio = new - old + else: + log_ratio = new - new.detach() + + if importance_sampling_level == "token": + log_importance_weights = log_ratio + elif importance_sampling_level == "sequence": + log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0) + log_importance_weights = log_importance_weights.unsqueeze(-1) + else: + raise ValueError( + f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' " + "and 'sequence'." + ) + + coef_1 = torch.exp(log_importance_weights) + + # Reverse KL: low-variance low-bias estimator as used in the GRPO paper. + if beta != 0.0: + kl_i = torch.exp(ref - new) - (ref - new) - 1.0 + + else: + # Zeros with the correct shape. + if importance_sampling_level == "sequence": + kl_i = new.new_zeros(new.size(0), 1) + else: + kl_i = torch.zeros_like(new) + + if loss_type == "cispo": + clamped_ratios = torch.clamp(coef_1, max=epsilon_high).detach() + loss_i = -clamped_ratios * advantages * new + elif loss_type in ["grpo", "bnpo", "dr_grpo", "dapo"]: + coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high) + + if delta is not None: + loss_1 = torch.clamp(coef_1, max=delta) * advantages + else: + loss_1 = coef_1 * advantages + pass + loss_2 = coef_2 * advantages + loss_i = -torch.min(loss_1, loss_2) + elif loss_type == "sapo": + if get_sapo_token_loss is None: + raise Exception(f"sapo is only available in TRL 0.26.0+") + loss_i = torch.empty_like(coef_1) + positive_advantages_mask = advantages.repeat([1, coef_1.shape[1]]) > 0 + # With n_chunks some tensors may be empty; guard the indexing. + if coef_1[positive_advantages_mask].numel() != 0: + loss_i[positive_advantages_mask] = get_sapo_token_loss( + coef_1[positive_advantages_mask], sapo_temperature_pos + ) + if coef_1[~positive_advantages_mask].numel() != 0: + loss_i[~positive_advantages_mask] = get_sapo_token_loss( + coef_1[~positive_advantages_mask], sapo_temperature_neg + ) + loss_i = -loss_i * advantages + elif loss_type == "vespo": + if get_gamma_weights is None: + raise Exception("vespo is only available in TRL 0.26.0+") + phi_seq = get_gamma_weights( + advantages=advantages, + log_ratio_per_token=log_ratio, + mask=mask, + importance_sampling_ratio=kwargs.get("importance_sampling_ratio"), + k_pos=vespo_k_pos, + lambda_pos=vespo_lambda_pos, + k_neg=vespo_k_neg, + lambda_neg=vespo_lambda_neg, + ) + loss_i = -phi_seq * advantages * new + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + if off_policy_mask_threshold is not None: + loss_i = loss_i * off_policy_mask + + if use_vllm and sampling_per_token_logps is not None: + loss_i = loss_i * importance_sampling_ratio + # delta for the metric. + with torch.no_grad(): + delta = torch.abs(old - sampling_per_token_logps) + delta = delta * mask + flat_is_ratio = importance_sampling_ratio * mask + else: + delta = torch.tensor([]).detach() + flat_is_ratio = torch.tensor([]).detach() + if beta != 0.0: + loss_i = loss_i + beta * kl_i + + mask = mask.to(torch.float32) + n_mask_per_reward = mask.sum(1) + + # https://github.com/huggingface/trl/blob/e8b8499f1f8d76838155b515e414ee98f757d6d5/trl/trainer/grpo_trainer.py#L1624 + if loss_type in ["grpo", "sapo"]: + loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() + loss = loss / current_gradient_accumulation_steps + elif loss_type == "bnpo": + loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0) + loss = loss / current_gradient_accumulation_steps + elif loss_type == "dr_grpo": + loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length) + loss = loss / current_gradient_accumulation_steps + elif loss_type in ["cispo", "dapo", "vespo"]: + normalizer = num_items_in_batch/ num_processes + loss = (loss_i * mask).sum() / normalizer + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + # Folded metrics. + def masked_batch_mean(x): + with torch.inference_mode(): + completion_length = n_mask_per_reward.mean() + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return completion_length, x.mean() + else: + mean_kl_per_reward = (x * mask).sum(1) / n_mask_per_reward + mean_kl = mean_kl_per_reward.mean() + return completion_length, mean_kl + completion_length, mean_kl = masked_batch_mean(kl_i) + return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, mask + +def grpo_update_SamplingParams(SamplingParams, generation_kwargs, vllm_sampling_params = None): + good_sampling_params_keys = inspect.signature(SamplingParams).parameters.keys() + + new_generation_kwargs = {} + for key in generation_kwargs.keys(): + if key in good_sampling_params_keys: + new_generation_kwargs[key] = generation_kwargs[key] + generation_kwargs = new_generation_kwargs + + if vllm_sampling_params is not None: + for key in good_sampling_params_keys: + if hasattr(vllm_sampling_params, key): + overwrited_key = getattr(vllm_sampling_params, key) + if overwrited_key is not None and (type(overwrited_key) in (list, tuple,) and len(overwrited_key) != 0): + generation_kwargs[key] = overwrited_key + return generation_kwargs + +def _get_inference_mode_context_manager(model: torch.nn.Module): + """ + If the state dict was quantized using torchao, we will run into + the following error when calling ops like aten.t() in inference mode. + This is a bug in PyTorch that affects all tensor subclasses. + + Cannot set version_counter for inference tensor + + For now, we work around this issue by using `torch.no_grad()` in this case. + See https://github.com/pytorch/pytorch/issues/164872 for more details. + Otherwise, just return `torch.inference_mode()`. + """ + torchao_config = getattr(model, "torchao_config", None) + if torchao_config is not None and torchao_config.qat_scheme is None: + return torch.no_grad() + else: + return torch.inference_mode() +@dataclass +class UnslothGRPOConfig(GRPOConfig): + """ + + Configuration class for the [`GRPOTrainer`]. + + This class includes only the parameters that are specific to GRPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`str`, `dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`GRPOTrainer`] is provided as a string. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to + [`~transformers.AutoModelForCausalLM.from_pretrained`] and + [`~transformers.AutoProcessor.from_pretrained`]. Also applied to reward-model and reward-tokenizer loads. + router_aux_loss_coef (`float`, *optional*, defaults to `0.001`): + Coefficient of the load-balancing auxiliary loss. Only has an effect when training a Mixture-of-Experts + (MoE) model; for other models it does nothing. The auxiliary loss is added to the training loss with this + weight. Set to `0.0` to disable it. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + cast_lm_head_to_fp32 (`bool`, *optional*, defaults to `False`): + Whether to cast the language modeling head of the policy and reference models to float32. As recommended by + the [ScaleRL](https://huggingface.co/papers/2510.13786) recipe. This flag is only supported when the model + has untied word embedding and language modeling head layers i.e. `tie_word_embeddings` in the model config + is False. + + > Parameters that control the data preprocessing + + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + num_generations (`int`, *optional*, defaults to `8`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + num_generations_eval (`int` or `None`, *optional*): + Number of generations to sample during evaluation. This allows using fewer generations during evaluation to + save computation. If `None`, uses the value of `num_generations`. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + pad_to_multiple_of (`int`, *optional*): + If set, the prompts ids and completions ids will be padded to a multiple of this value. + + > Parameters that control generation + + generation_batch_size (`int`, *optional*): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one + generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. + steps_per_generation (`int`, *optional*): + Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive + with `generation_batch_size`. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int`, *optional*, defaults to `0`): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + generation_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or + `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the + generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict + with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. + chat_template_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + cache_implementation (`str`, *optional*): + Implementation of the cache method for faster generation when `use_vllm` is set to `False`. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"colocate"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + vllm_structured_outputs_regex (`str`, *optional*): + Regex for vLLM structured outputs. If `None` (default), structured outputs is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + vllm_group_port (`int`, *optional*, defaults to `51216`): + Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port + is occupied, there is no need to change it. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_max_model_length (`int`, *optional*): + Context window for vLLM. Set it to at least the maximum prompt length in the dataset plus + `max_completion_length`; if omitted, it is inferred from the model config. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but + waking the engine adds host–device transfer latency. + + > Parameters that control generation acceleration powered by transformers continuous batching + + use_transformers_continuous_batching (`bool`, *optional*, defaults to `False`): + Whether to use transformers' continuous batching engine for generating completions. Requires + `transformers>=5.8.0`. + transformers_continuous_batching_config (`dict`, *optional*): + Keyword arguments for [`~transformers.generation.ContinuousBatchingConfig`]. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.0`): + KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and improving + training speed. [DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement + learning](https://huggingface.co/papers/2501.12948) use a value of `0.001`. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + delta (`float`, *optional*): + Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard + GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in + the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291). + epsilon_high (`float`, *optional*): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + When used with `loss_type='cispo'`, this corresponds to the ε_max param specified in the [ScaleRL + paper](https://huggingface.co/papers/2510.13786) and the recommended value is `5.0`. + sapo_temperature_neg (`float`, *optional*, defaults to `1.05`): + Temperature for tokens with non-positive advantage scores used in the `sapo` loss function. This parameter + is introduced in the [Soft Adaptive Policy Optimization paper](https://huggingface.co/papers/2511.20347). + sapo_temperature_pos (`float`, *optional*, defaults to `1.0`): + Temperature for tokens with positive advantage scores used in the `sapo` loss function. This parameter is + introduced in the [Soft Adaptive Policy Optimization paper](https://huggingface.co/papers/2511.20347). + vespo_k_pos (`float`, *optional*, defaults to `2.0`): + k parameter for positive advantages, it is the power exponent in the VESPO loss. Controls how aggressively + we down-weight samples with low importance weights (when the importance sampling ratio < 1). + vespo_lambda_pos (`float`, *optional*, defaults to `3.0`): + lambda parameter for positive advantages, it is the decay factor in the VESPO loss. Controls how + aggressively we down-weight samples with high importance weights (when the importance sampling ratio > 1). + vespo_k_neg (`float`, *optional*, defaults to `3.0`): + k parameter for negative advantages, it is the power exponent in the VESPO loss. Controls how aggressively + we down-weight samples with low importance weights (when the importance sampling ratio < 1). + vespo_lambda_neg (`float`, *optional*, defaults to `2.0`): + lambda parameter for negative advantages, it is the exponential decay factor in the VESPO loss. Controls + how aggressively we down-weight samples with high importance weights (when the importance sampling ratio > + 1). + importance_sampling_level (`str`, *optional*, defaults to `"token"`): + Controls whether importance sampling ratios are computed at the `"token"` or `"sequence"` level. `"token"` + keeps the raw per-token log-probability ratios (one weight per token). `"sequence"` averages the + log-probability ratios across valid tokens to produce a single ratio per sequence. The [GSPO + paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more + stable training and better alignment with sequence-level rewards. + reward_weights (`list[float]`, *optional*): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + multi_objective_aggregation (`str`, *optional*, defaults to `"sum_then_normalize"`): + Method to aggregate multiple reward functions. Supported values are: + + - `"sum_then_normalize"` (default): First sums the weighted rewards from each reward function, then applies + reward scaling/normalization as specified by `scale_rewards` (see `scale_rewards` for details). + - `"normalize_then_sum"`: First normalizes/scales each reward function across generations (within each + group), then sums the normalized rewards using the specified weights. The aggregated reward is then + normalized at the batch level when forming advantages. This is the suggested approach from the paper + [GDPO: Group reward-Decoupled Normalization Policy Optimization for Multi-reward RL + Optimization](https://huggingface.co/papers/2601.05242). + scale_rewards (`str` or `bool`, *optional*, defaults to `"group"`): + Specifies the scaling strategy for rewards. Supported values are: + + - `True` or `"group"` (default): rewards are scaled by the standard deviation within each group, ensuring + unit variance within a group. + - `"batch"`: rewards are scaled by the standard deviation across the entire batch, as recommended in the + [PPO Lite paper](https://huggingface.co/papers/2508.08221). + - `False` or `"none"`: no scaling is applied. The [Dr. GRPO + paper](https://huggingface.co/papers/2503.20783) recommends not scaling rewards, as scaling by the + standard deviation introduces a question-level difficulty bias. + loss_type (`str`, *optional*, defaults to `"dapo"`): + Specifies the loss formulation to use. Supported values are: + + - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to + length bias—this approach tends to prefer shorter completions with positive advantages and longer ones + with negative advantages. + - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was + introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias. + The value of the constant corresponds to `max_completion_length`. + - `"dapo"` (default): Aggregates token-level losses by normalizing with the number of active token in the + global accumulated batch. This method was introduced in the [DAPO + paper](https://huggingface.co/papers/2503.14476) to eliminate length bias. + - `"bnpo"`: Aggregates token-level losses by normalizing with the number of active token in the local + batch. Note that normalization is performed over the local batch only, so results may slightly vary + depending on the local batch size, despite a constant effective batch size. When using + `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. + - `"cispo"`: Clips the importance sampling weights instead of the advantage scaled importance weights. The + clipped weights are then multiplied with the advantages and policy model's log probs. Individual token + losses are aggregated by normalizing with the number of active tokens in the global accumulated batch. + This method was introduced in the [MiniMax-M1 paper](https://huggingface.co/papers/2506.13585). + - `"sapo"`: Soft Adaptive Policy Optimization loss, as introduced in the [Soft Adaptive Policy Optimization + paper](https://huggingface.co/papers/2511.20347). Replaces hard clipping with a smooth, + temperature-controlled gate that adaptively attenuates off-policy updates while preserving useful + learning signals. + - `"luspo"`: Length-Unbiased Sequence Policy Optimization loss. A sequence-level loss that scales each + sequence's loss by its length. This is a modification of GSPO and requires + `importance_sampling_level="sequence"`. Introduced in the [LUSPO + paper](https://huggingface.co/papers/2602.05261). + - `"vespo"`: Variational Sequence-Level Soft Policy Optimization. Replaces hard clipping with a smooth, + asymmetric Gamma weighting function applied directly to sequence-level importance weights. Introduced in + the [VESPO paper](https://huggingface.co/papers/2602.10693). + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + top_entropy_quantile (`float`, *optional*, defaults to `1.0`): + ρ parameter from [Beyond the 80/20 Rule](https://huggingface.co/papers/2506.01939). Keeps in the policy + loss term only the top-ρ quantile of tokens by entropy of the probability distribution at each sequence + position, improving results. Range: `[0.0-1.0]`. A value of `0.0` masks all but the highest entropy token; + `1.0` keeps all tokens. The paper recommends a value of `0.2`. If used with + `mask_truncated_completions=True`, only tokens from non-truncated completions are considered. + max_tool_calling_iterations (`int`, *optional*): + Maximum number of tool-calling turns when training an agent. If `None`, there is no limit and generation + stops when the model generates a response turn with no tool calls or when the total response length reaches + `max_model_length`. + vllm_importance_sampling_correction (`bool`, *optional*, defaults to `True`): + Whether to apply Importance Sampling (IS) to correct for the mismatch between vLLM completion logprobs and + recomputed training logprobs. If set to `False`, no IS is applied regardless of + `vllm_importance_sampling_mode`. When `True`, the selected mode determines how the IS ratios are computed + and constrained. + vllm_importance_sampling_mode (`str`, *optional*, defaults to `"sequence_mask"`): + Specifies how Importance Sampling is performed when `vllm_importance_sampling_correction=True`. Possible + values are: + + - `"token_truncate"`: Token-level truncated IS (default). Per-token ratios are clipped to + [C_min, C_max]. + - `"token_mask"`: Token-level masked IS. Per-token ratios outside [C_min, C_max] are set to zero. + - `"sequence_truncate"`: Sequence-level truncated IS. A single sequence ratio is clipped to + [C_min, C_max] and applied to all tokens in the sequence. + - `"sequence_mask"`: Sequence-level masked IS. Sequences with ratios outside [C_min, C_max] are masked + out. + vllm_importance_sampling_clip_max (`float`, *optional*, defaults to `3.0`): + Importance sampling upper bound C_max used by `vllm_importance_sampling_mode`. For `*_truncate` modes, + importance ratios are clipped from above at C_max. For `*_mask` modes, ratios larger than C_max are set to + zero. + vllm_importance_sampling_clip_min (`float`, *optional*): + Importance sampling lower bound C_min used by `vllm_importance_sampling_mode`. For `*_truncate` modes, + ratios are clipped from below at C_min. For `*_mask` modes, ratios below C_min are set to zero. To strictly + mask ratios below C_min without upper bound, set `vllm_importance_sampling_clip_max=None`. + off_policy_mask_threshold (`float`, *optional*): + Threshold for off-policy sequence masking. If `None`, off-policy sequence masking is disabled. When set, + sequences with negative advantages and high KL divergence are masked out to stabilize training. This + parameter corresponds to the `delta` threshold in Equation 9 of the [DeepSeek-V3.2 + paper](https://huggingface.co/papers/2512.02556). It expects a positive value (e.g., 0.5). + use_bias_correction_kl (`bool`, *optional*, defaults to `False`): + Whether to use the unbiased KL divergence estimator with importance sampling correction. This corrects the + KL divergence estimate by multiplying it with the importance sampling ratio. This is described in the + [DeepSeek-V3.2 paper](https://huggingface.co/papers/2512.02556). + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or + `trackio`. + num_completions_to_print (`int`, *optional*): + Number of completions to print with `rich`. If `None`, all completions are logged. + log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts. If `True`, only unique prompts are logged. If `False`, all prompts are + logged. + log_completions_hub_repo (`str`, *optional*): + Hugging Face Hub repository to save the completions. Should be a complete repository name like + `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` in which case the repository will be + created in the currently-logged-in Hugging Face user's namespace. Note that this repository will be public + unless you set `hub_private_repo=True` or your organization's default is to create private repositories." + + > Deprecated parameters + + use_transformers_paged: + + + + Parameter `use_transformers_paged` is deprecated and will be removed in version v2.0.0. Use + `use_transformers_continuous_batching` instead. + + + + vllm_importance_sampling_cap: + + + + Parameter `vllm_importance_sampling_cap` is deprecated and will be removed in v2.0.0. Use + `vllm_importance_sampling_clip_max` instead. + + + + > [!NOTE] + > These parameters have default values different from [`~transformers.TrainingArguments`]: + > - `logging_steps`: Defaults to `10` instead of `500`. + > - `gradient_checkpointing`: Defaults to `True` instead of `False`. + > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. + > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. + + """ + vllm_sampling_params: Optional[Any] = field( + default = None, + metadata = {'help': 'vLLM SamplingParams'}, + ) + unsloth_num_chunks : Optional[int] = field( + default = -1, + metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, + ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {'help': 'Multiplier for chunked logit computations.'}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, + ) + + def __init__( + self, + output_dir = None, + per_device_train_batch_size = 4, + num_train_epochs = 3.0, + max_steps = -1, + learning_rate = 5e-05, + lr_scheduler_type = 'linear', + lr_scheduler_kwargs = None, + warmup_steps = 0.1, + optim = 'adamw_8bit', + optim_args = None, + weight_decay = 0.001, + adam_beta1 = 0.9, + adam_beta2 = 0.999, + adam_epsilon = 1e-08, + optim_target_modules = None, + gradient_accumulation_steps = 2, + average_tokens_across_devices = True, + max_grad_norm = 1.0, + label_smoothing_factor = 0.0, + bf16 = False, + fp16 = False, + bf16_full_eval = False, + fp16_full_eval = False, + tf32 = None, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = None, + torch_compile = False, + torch_compile_backend = None, + torch_compile_mode = None, + use_liger_kernel = False, + liger_kernel_config = None, + use_cache = False, + neftune_noise_alpha = None, + torch_empty_cache_steps = 250, + auto_find_batch_size = False, + logging_strategy = 'steps', + logging_steps = 1, + logging_first_step = False, + log_on_each_node = True, + logging_nan_inf_filter = False, + include_num_input_tokens_seen = False, + log_level = 'passive', + log_level_replica = 'warning', + disable_tqdm = None, + report_to = 'none', + run_name = None, + project = 'huggingface', + trackio_space_id = 'trackio', + eval_strategy = 'no', + eval_steps = None, + eval_delay = 0, + per_device_eval_batch_size = 4, + prediction_loss_only = False, + eval_on_start = False, + eval_do_concat_batches = True, + eval_use_gather_object = False, + eval_accumulation_steps = 2, + batch_eval_metrics = False, + save_only_model = False, + save_strategy = 'steps', + save_steps = 500, + save_on_each_node = False, + save_total_limit = None, + enable_jit_checkpoint = False, + push_to_hub = False, + hub_token = None, + hub_private_repo = None, + hub_model_id = None, + hub_strategy = 'every_save', + hub_always_push = False, + hub_revision = None, + load_best_model_at_end = False, + metric_for_best_model = None, + greater_is_better = None, + ignore_data_skip = False, + restore_callback_states_from_checkpoint = False, + full_determinism = False, + seed = 3407, + data_seed = 3407, + use_cpu = False, + accelerator_config = None, + parallelism_config = None, + dataloader_drop_last = False, + dataloader_num_workers = 0, + dataloader_pin_memory = True, + dataloader_persistent_workers = False, + dataloader_prefetch_factor = None, + remove_unused_columns = False, + label_names = None, + train_sampling_strategy = 'random', + length_column_name = 'length', + ddp_find_unused_parameters = None, + ddp_bucket_cap_mb = None, + ddp_broadcast_buffers = None, + ddp_backend = None, + ddp_timeout = 1800, + fsdp = None, + fsdp_config = None, + deepspeed = None, + debug = '', + skip_memory_metrics = True, + do_train = False, + do_eval = False, + do_predict = False, + resume_from_checkpoint = None, + warmup_ratio = None, + logging_dir = None, + local_rank = -1, + model_init_kwargs = None, + trust_remote_code = False, + router_aux_loss_coef = 0.001, + disable_dropout = False, + cast_lm_head_to_fp32 = False, + num_generations = 8, + num_generations_eval = None, + max_completion_length = 256, + ds3_gather_for_generation = True, + shuffle_dataset = True, + pad_to_multiple_of = None, + generation_batch_size = None, + steps_per_generation = None, + temperature = 1.0, + top_p = 1.0, + top_k = None, + min_p = None, + generation_kwargs = {}, + chat_template_kwargs = None, + repetition_penalty = 1.0, + cache_implementation = None, + use_vllm = False, + vllm_mode = 'colocate', + vllm_model_impl = 'vllm', + vllm_enable_sleep_mode = False, + vllm_structured_outputs_regex = None, + vllm_server_base_url = None, + vllm_server_host = '0.0.0.0', + vllm_server_port = 8000, + vllm_server_timeout = 240.0, + vllm_group_port = 51216, + vllm_gpu_memory_utilization = 0.3, + vllm_max_model_length = None, + vllm_tensor_parallel_size = 1, + beta = 0.001, + num_iterations = 1, + epsilon = 0.2, + delta = None, + epsilon_high = None, + sapo_temperature_neg = 1.05, + sapo_temperature_pos = 1.0, + vespo_k_pos = 2.0, + vespo_lambda_pos = 3.0, + vespo_k_neg = 3.0, + vespo_lambda_neg = 2.0, + importance_sampling_level = 'token', + reward_weights = None, + multi_objective_aggregation = 'sum_then_normalize', + scale_rewards = 'group', + loss_type = 'bnpo', + mask_truncated_completions = False, + sync_ref_model = False, + ref_model_mixup_alpha = 0.6, + ref_model_sync_steps = 512, + top_entropy_quantile = 1.0, + max_tool_calling_iterations = None, + vllm_importance_sampling_correction = False, + vllm_importance_sampling_mode = 'sequence_mask', + vllm_importance_sampling_clip_max = 3.0, + vllm_importance_sampling_clip_min = None, + off_policy_mask_threshold = None, + use_bias_correction_kl = False, + log_completions = False, + num_completions_to_print = None, + log_unique_prompts = False, + log_completions_hub_repo = None, + use_transformers_continuous_batching = False, + transformers_continuous_batching_config = None, + use_transformers_paged = False, + vllm_importance_sampling_cap = None, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, + + **kwargs, + ): + if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') + if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') + if num_train_epochs is None: + num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override + if output_dir is None and save_strategy == 'steps' and save_steps == 500: + output_dir = 'unsloth_training_checkpoints' + save_strategy = 'no' + if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': + from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION + if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: + from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE + pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE + + if loss_type.lower() == 'dr_grpo': + loss_type = 'dr_grpo' + elif loss_type.lower() == 'dapo': + loss_type = 'dapo' + if loss_type.lower() == 'dr_grpo': + if scale_rewards == None: + scale_rewards = True + elif scale_rewards == True: + print('Unsloth: The Dr GRPO paper recommends setting `scale_rewards` to False! Will override. Set it to `None` to force False.') + scale_rewards = False + elif loss_type.lower() == 'dapo': + if mask_truncated_completions != True: + print('Unsloth: The DAPO paper recommends `mask_truncated_completions = True` - we will set it.') + if epsilon_high != 0.28: + print('Unsloth: The DAPO paper recommends `epsilon_high = 0.28` - we will set it.') + if beta != 0.0: + print(f'[WARNING] Unsloth: The DAPO paper recommends setting `beta = 0.0` to remove the KL term - You have set it to {beta}.') + mask_truncated_completions = True + epsilon_high = 0.28 + + if steps_per_generation is None and generation_batch_size is None: + ga = gradient_accumulation_steps + world_size = int(os.environ.get('WORLD_SIZE', '1')) + if (ga * world_size * per_device_train_batch_size) % num_generations != 0: + print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) + per_device_train_batch_size = num_generations + + if temperature <= 0: + raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') + + if use_vllm and (top_k is None or top_k == 0): top_k = -1 + + super().__init__( + output_dir = output_dir, + per_device_train_batch_size = per_device_train_batch_size, + num_train_epochs = num_train_epochs, + max_steps = max_steps, + learning_rate = learning_rate, + lr_scheduler_type = lr_scheduler_type, + lr_scheduler_kwargs = lr_scheduler_kwargs, + warmup_steps = warmup_steps, + optim = optim, + optim_args = optim_args, + weight_decay = weight_decay, + adam_beta1 = adam_beta1, + adam_beta2 = adam_beta2, + adam_epsilon = adam_epsilon, + optim_target_modules = optim_target_modules, + gradient_accumulation_steps = gradient_accumulation_steps, + average_tokens_across_devices = average_tokens_across_devices, + max_grad_norm = max_grad_norm, + label_smoothing_factor = label_smoothing_factor, + bf16 = bf16, + fp16 = fp16, + bf16_full_eval = bf16_full_eval, + fp16_full_eval = fp16_full_eval, + tf32 = tf32, + gradient_checkpointing = gradient_checkpointing, + gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, + torch_compile = torch_compile, + torch_compile_backend = torch_compile_backend, + torch_compile_mode = torch_compile_mode, + use_liger_kernel = use_liger_kernel, + liger_kernel_config = liger_kernel_config, + use_cache = use_cache, + neftune_noise_alpha = neftune_noise_alpha, + torch_empty_cache_steps = torch_empty_cache_steps, + auto_find_batch_size = auto_find_batch_size, + logging_strategy = logging_strategy, + logging_steps = logging_steps, + logging_first_step = logging_first_step, + log_on_each_node = log_on_each_node, + logging_nan_inf_filter = logging_nan_inf_filter, + include_num_input_tokens_seen = include_num_input_tokens_seen, + log_level = log_level, + log_level_replica = log_level_replica, + disable_tqdm = disable_tqdm, + report_to = report_to, + run_name = run_name, + project = project, + trackio_space_id = trackio_space_id, + eval_strategy = eval_strategy, + eval_steps = eval_steps, + eval_delay = eval_delay, + per_device_eval_batch_size = per_device_eval_batch_size, + prediction_loss_only = prediction_loss_only, + eval_on_start = eval_on_start, + eval_do_concat_batches = eval_do_concat_batches, + eval_use_gather_object = eval_use_gather_object, + eval_accumulation_steps = eval_accumulation_steps, + batch_eval_metrics = batch_eval_metrics, + save_only_model = save_only_model, + save_strategy = save_strategy, + save_steps = save_steps, + save_on_each_node = save_on_each_node, + save_total_limit = save_total_limit, + enable_jit_checkpoint = enable_jit_checkpoint, + push_to_hub = push_to_hub, + hub_token = hub_token, + hub_private_repo = hub_private_repo, + hub_model_id = hub_model_id, + hub_strategy = hub_strategy, + hub_always_push = hub_always_push, + hub_revision = hub_revision, + load_best_model_at_end = load_best_model_at_end, + metric_for_best_model = metric_for_best_model, + greater_is_better = greater_is_better, + ignore_data_skip = ignore_data_skip, + restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, + full_determinism = full_determinism, + seed = seed, + data_seed = data_seed, + use_cpu = use_cpu, + accelerator_config = accelerator_config, + parallelism_config = parallelism_config, + dataloader_drop_last = dataloader_drop_last, + dataloader_num_workers = dataloader_num_workers, + dataloader_pin_memory = dataloader_pin_memory, + dataloader_persistent_workers = dataloader_persistent_workers, + dataloader_prefetch_factor = dataloader_prefetch_factor, + remove_unused_columns = remove_unused_columns, + label_names = label_names, + train_sampling_strategy = train_sampling_strategy, + length_column_name = length_column_name, + ddp_find_unused_parameters = ddp_find_unused_parameters, + ddp_bucket_cap_mb = ddp_bucket_cap_mb, + ddp_broadcast_buffers = ddp_broadcast_buffers, + ddp_backend = ddp_backend, + ddp_timeout = ddp_timeout, + fsdp = fsdp, + fsdp_config = fsdp_config, + deepspeed = deepspeed, + debug = debug, + skip_memory_metrics = skip_memory_metrics, + do_train = do_train, + do_eval = do_eval, + do_predict = do_predict, + resume_from_checkpoint = resume_from_checkpoint, + warmup_ratio = warmup_ratio, + logging_dir = logging_dir, + local_rank = local_rank, + model_init_kwargs = model_init_kwargs, + trust_remote_code = trust_remote_code, + router_aux_loss_coef = router_aux_loss_coef, + disable_dropout = disable_dropout, + cast_lm_head_to_fp32 = cast_lm_head_to_fp32, + num_generations = num_generations, + num_generations_eval = num_generations_eval, + max_completion_length = max_completion_length, + ds3_gather_for_generation = ds3_gather_for_generation, + shuffle_dataset = shuffle_dataset, + pad_to_multiple_of = pad_to_multiple_of, + generation_batch_size = generation_batch_size, + steps_per_generation = steps_per_generation, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + generation_kwargs = generation_kwargs, + chat_template_kwargs = chat_template_kwargs, + repetition_penalty = repetition_penalty, + cache_implementation = cache_implementation, + use_vllm = use_vllm, + vllm_mode = vllm_mode, + vllm_model_impl = vllm_model_impl, + vllm_enable_sleep_mode = vllm_enable_sleep_mode, + vllm_structured_outputs_regex = vllm_structured_outputs_regex, + vllm_server_base_url = vllm_server_base_url, + vllm_server_host = vllm_server_host, + vllm_server_port = vllm_server_port, + vllm_server_timeout = vllm_server_timeout, + vllm_group_port = vllm_group_port, + vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, + vllm_max_model_length = vllm_max_model_length, + vllm_tensor_parallel_size = vllm_tensor_parallel_size, + beta = beta, + num_iterations = num_iterations, + epsilon = epsilon, + delta = delta, + epsilon_high = epsilon_high, + sapo_temperature_neg = sapo_temperature_neg, + sapo_temperature_pos = sapo_temperature_pos, + vespo_k_pos = vespo_k_pos, + vespo_lambda_pos = vespo_lambda_pos, + vespo_k_neg = vespo_k_neg, + vespo_lambda_neg = vespo_lambda_neg, + importance_sampling_level = importance_sampling_level, + reward_weights = reward_weights, + multi_objective_aggregation = multi_objective_aggregation, + scale_rewards = scale_rewards, + loss_type = loss_type, + mask_truncated_completions = mask_truncated_completions, + sync_ref_model = sync_ref_model, + ref_model_mixup_alpha = ref_model_mixup_alpha, + ref_model_sync_steps = ref_model_sync_steps, + top_entropy_quantile = top_entropy_quantile, + max_tool_calling_iterations = max_tool_calling_iterations, + vllm_importance_sampling_correction = vllm_importance_sampling_correction, + vllm_importance_sampling_mode = vllm_importance_sampling_mode, + vllm_importance_sampling_clip_max = vllm_importance_sampling_clip_max, + vllm_importance_sampling_clip_min = vllm_importance_sampling_clip_min, + off_policy_mask_threshold = off_policy_mask_threshold, + use_bias_correction_kl = use_bias_correction_kl, + log_completions = log_completions, + num_completions_to_print = num_completions_to_print, + log_unique_prompts = log_unique_prompts, + log_completions_hub_repo = log_completions_hub_repo, + use_transformers_continuous_batching = use_transformers_continuous_batching, + transformers_continuous_batching_config = transformers_continuous_batching_config, + use_transformers_paged = use_transformers_paged, + vllm_importance_sampling_cap = vllm_importance_sampling_cap,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier + + # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+ + if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: + if 'use_reentrant' in self.gradient_checkpointing_kwargs: + del self.gradient_checkpointing_kwargs['use_reentrant'] + +pass + +class _UnslothGRPOTrainer(_BaseTrainer): + """""" + + _tag_names = ["trl", "grpo"] + _name = "GRPO" + _paper = { + "title": "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models", + "id": "2402.03300", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{shao2024deepseekmath, + title = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}}, + author = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo}, + year = 2024, + eprint = {arXiv:2402.03300}, + }"""), + } + + def __init__( + self, + model: "str | PreTrainedModel | PeftModel", + reward_funcs: RewardFunc | list[RewardFunc], + args: GRPOConfig | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + peft_config: "PeftConfig | None" = None, + tools: list[Callable] | None = None, + rollout_func: RolloutFunc | None = None, + environment_factory: EnvironmentFactory | None = None, + ): + + if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): + if (getattr(args, 'use_vllm', False) == False): + args.use_vllm = True + args.vllm_mode='colocate' + _unsloth_esm = getattr(getattr(getattr(getattr(model.vllm_engine, 'llm_engine', None), 'vllm_config', None), 'model_config', None), 'enable_sleep_mode', None) + if (_unsloth_esm if _unsloth_esm is not None else os.environ.get('UNSLOTH_VLLM_STANDBY', '0') != '0'): + args.vllm_enable_sleep_mode=True + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = GRPOConfig(f"{model_name}-GRPO") + + # Model + if isinstance(model, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + model = create_model_from_path(model, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do + _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + + # Some models [SmolVLM/Idefics3] don't support `logits_to_keep` argument and error out if we pass it + # Inspect the forward method before we wrap the model with PEFT + self.model_kwarg_keys = ( + inspect.signature(model.forward).parameters.keys() + if not hasattr(model, "get_base_model") + else inspect.signature(model.get_base_model().forward).parameters.keys() + ) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained( + get_config_model_id(model.config), + truncation_side="left", + padding_side="left", + trust_remote_code=args.trust_remote_code, + ) + + if args.use_transformers_continuous_batching and isinstance(processing_class, ProcessorMixin): + raise ValueError( + "`use_transformers_continuous_batching` does not support multimodal models. Use `use_vllm` instead." + ) + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + self._tokenizer = processing_class.tokenizer + self._is_vlm = True + elif isinstance(processing_class, PreTrainedTokenizerBase): + self._tokenizer = processing_class + self._is_vlm = False + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if self._tokenizer.pad_token is None: + self._tokenizer.pad_token = self._tokenizer.eos_token + + # Resolve vision placeholder token IDs once. Used by the forward pass to rebuild mm_token_type_ids + # when tool responses inject images into the completion [see _generate forward_kwargs block]. + self._image_pad_token_id = None + self._video_pad_token_id = None + if self._is_vlm: + for candidate in ("<|image_pad|>", "<|image|>"): + tid = self._tokenizer.convert_tokens_to_ids(candidate) + if tid != self._tokenizer.unk_token_id: + self._image_pad_token_id = tid + break + tid = self._tokenizer.convert_tokens_to_ids("<|video_pad|>") + if tid != self._tokenizer.unk_token_id: + self._video_pad_token_id = tid + + # PEFT + if False: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if is_peft_model(model): + raise ValueError( + "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " + "and unload the existing adapter, save the resulting base model, and then pass that base model along " + "with the new `peft_config` to the trainer." + ) + # Create PEFT model + # ZeRO-3 + PEFT for non-quantized models: + # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. + # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, + # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. + # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. + # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for + # non-quantized bf16 training. + # - See: + # - TRL issue: https://github.com/huggingface/trl/issues/6089 + # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 + # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg + get_peft_model_kwargs = {} + if ( + args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and not _is_quantized_model + and Version(peft.__version__) >= Version("0.12.0") + ): + get_peft_model_kwargs["autocast_adapter_dtype"] = False + model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + + elif is_peft_model(model) and args.beta != 0.0: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during GRPO training. PEFT only + # supports one adapter per model when the LoRA config uses `target_parameters` [see peft#3340], so in that + # case we skip the "ref" adapter and compute the reference log probs with adapters disabled, i.e. with the + # base model. + default_config = model.peft_config["default"] + if isinstance(default_config, LoraConfig) and default_config.target_parameters: + logger.warning( + "PEFT can't add a frozen reference adapter alongside one that uses `target_parameters` " + "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " + "If you wrapped the model only to apply LoRA, pass a `peft_config` to the trainer instead; if you " + "wrapped it deliberately (pretrained adapter or custom init), note that the base model matches " + "your adapter only when it's freshly zero-initialized. If it is, this warning is safe to ignore." + ) + else: + model.add_adapter("ref", default_config) + for name, param in model.named_parameters(): + if ".default." in name: + ref_name = name.replace(".default.", ".ref.") + ref_param = model.get_parameter(ref_name) + ref_param.data.copy_(param.data) + + # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally + # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 + if is_peft_model(model) and args.gradient_checkpointing: + model.enable_input_require_grads() + + # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the + # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by + # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for + # quantized models. See: https://github.com/huggingface/peft/issues/2889 + if _is_quantized_model: + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.bfloat16) + + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError( + f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " + f"reward functions ({len(reward_funcs)})." + ) + + for i, (reward_processing_class, reward_func) in enumerate( + zip(reward_processing_classes, reward_funcs, strict=True) + ): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained( + get_config_model_id(reward_func.config), trust_remote_code=args.trust_remote_code + ) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + + self.reward_processing_classes = reward_processing_classes + + # Rollout function + if rollout_func is not None and os.environ.get("TRL_EXPERIMENTAL_SILENCE", "0") != "1": + warnings.warn( + "You are using 'rollout_func', which is an experimental feature. This API may change or be removed at " + "any time without prior notice. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1.", + UserWarning, + stacklevel=2, + ) + self.rollout_func = rollout_func + if environment_factory is not None and os.environ.get("TRL_EXPERIMENTAL_SILENCE", "0") != "1": + warnings.warn( + "You are using 'environment_factory', which is an experimental feature. This API may change or be " + "removed at any time without prior notice. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1.", + UserWarning, + stacklevel=2, + ) + + # Tools + if tools: + if not Version(transformers.__version__) >= Version("5.0.0"): + raise ImportError( + "Using tools with GRPOTrainer requires transformers version 5.0.0 or higher. Please upgrade " + "transformers with `pip install --upgrade transformers` to use this feature." + ) + if environment_factory: + if not Version(transformers.__version__) >= Version("5.2.0"): + raise ImportError( + "Using `environment_factory` with GRPOTrainer requires transformers version 5.2.0 or higher. " + "Please install transformers from the main branch with `pip install " + "git+https://github.com/huggingface/transformers.git@main` to use this feature." + ) + if tools or environment_factory: + if not is_jmespath_available(): + raise ImportError( + "Using tools with GRPOTrainer requires the jmespath library for response parsing. Please install " + "it with `pip install jmespath` to use this feature." + ) + if not supports_tool_calling(processing_class): + raise ValueError( + "The provided chat template does not support tool calling. The template must be able to render a " + "full tool-calling conversation (user -> assistant with tool_calls -> tool)." + ) + + # Create the environments and extract their methods to be used as tools. We create one environment per rollout + generation_batch_size = args.per_device_train_batch_size * args.steps_per_generation + if environment_factory is not None: + self.environments = [environment_factory() for _ in range(generation_batch_size)] + environment_methods = [[] for _ in range(generation_batch_size)] + for i, environment in enumerate(self.environments): + has_reset = False + for name, member in inspect.getmembers(environment, predicate=inspect.ismethod): + if name == "reset": + has_reset = True + elif not name.startswith("_"): + environment_methods[i].append(member) + if not has_reset: + raise ValueError( + "Each environment instance returned by `environment_factory` must define a callable `reset` " + ) + else: + self.environments = None + + tools = tools or [] + self._sync_tool_dicts = [{} for _ in range(generation_batch_size)] + self._async_tool_dicts = [{} for _ in range(generation_batch_size)] + for i in range(generation_batch_size): + for tool in tools + (environment_methods[i] if self.environments is not None else []): + if inspect.iscoroutinefunction(tool): + self._async_tool_dicts[i][tool.__name__] = tool + else: + self._sync_tool_dicts[i][tool.__name__] = tool + + self.tools = tools + (environment_methods[0] if self.environments is not None else []) + + # Check for async functions to start an event loop on a daemon thread + self._has_async_funcs = any(inspect.iscoroutinefunction(func) for func in self.reward_funcs + self.tools) + + if self._has_async_funcs: + self.async_loop_thread, self.async_loop, self.async_loop_ready_event = start_event_loop_in_daemon( + name="GRPOTrainer-AsyncLoop" + ) + # wait until the event loop is running in the daemon thread + self.async_loop_ready_event.wait() + atexit.register(shutdown_event_loop_in_daemon, self.async_loop_thread, self.async_loop) + + # At the time of initial implementation, most tokenizers do not have built-in support for response schemas. + # While waiting for broader adoption, we provide this utility function to manually set the response schema for + # known chat templates. `response_schema` lives on the [inner] tokenizer, since `parse_response` is a tokenizer + # method that reads `self.response_schema`. + if self.tools and getattr(self._tokenizer, "response_schema", None) is None: + processing_class = add_response_schema(processing_class) + # In multi-turn training, the chat template *must* be prefix-preserving. If the tokenizer's original template + # isn't, we replace it at initialization with a training-safe, prefix-preserving template. + if self.tools and not is_chat_template_prefix_preserving(processing_class): + self.chat_template = get_training_chat_template(processing_class) + else: + self.chat_template = None + + # Training arguments + self.max_completion_length = args.max_completion_length # = |o_i| in the GRPO paper + self.num_generations = args.num_generations # = G in the GRPO paper + self.max_tool_calling_iterations = args.max_tool_calling_iterations or sys.maxsize + self.num_generations_eval = args.num_generations_eval or self.num_generations + self.chat_template_kwargs = args.chat_template_kwargs or {} + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_transformers_continuous_batching = args.use_transformers_continuous_batching + if self.use_transformers_continuous_batching: + if not Version(transformers.__version__) >= Version("5.8.0"): + raise ImportError( + "Using `use_transformers_continuous_batching` requires transformers>=5.8.0. " + "Please upgrade with `pip install --upgrade transformers`." + ) + from transformers.generation import ContinuousBatchingConfig + + cb_kwargs = dict(args.transformers_continuous_batching_config or {}) + # The transformers default [0.9] leaves almost no VRAM for the training backward pass; + # use a training-aware default unless the user has set it explicitly. + cb_kwargs.setdefault("max_memory_percent", 0.5) + self.continuous_batching_config = ContinuousBatchingConfig(**cb_kwargs) + else: + self.continuous_batching_config = None + self.pad_to_multiple_of = args.pad_to_multiple_of + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction + self.vllm_importance_sampling_mode = args.vllm_importance_sampling_mode + self.vllm_importance_sampling_clip_max = args.vllm_importance_sampling_clip_max + self.vllm_importance_sampling_clip_min = args.vllm_importance_sampling_clip_min + self.use_liger_kernel = args.use_liger_kernel + self.loss_type = args.loss_type + self.multi_objective_aggregation = args.multi_objective_aggregation + + # MoE load-balancing auxiliary loss, applied to Mixture-of-Experts models [no effect otherwise] + text_config = model.config.get_text_config() + is_moe = getattr(text_config, "output_router_logits", None) is not None + self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0 + self.router_aux_loss_coef = args.router_aux_loss_coef + self.scale_rewards = args.scale_rewards + self.importance_sampling_level = args.importance_sampling_level + self.off_policy_mask_threshold = args.off_policy_mask_threshold + if self.use_liger_kernel and self.off_policy_mask_threshold is not None: + raise ValueError("Liger kernel does not support off-policy sequence masking yet.") + if self.use_liger_kernel and is_peft_model(model): + # The Liger fused GRPO loss multiplies the hidden states by `lm_head.weight` directly. When the LM head is + # targeted by a PEFT adapter [`"lm_head"` in `target_modules`], `lm_head.weight` is the frozen base weight + # and the trainable adapter parameters live in separate submodules that Liger never sees. The head adapter + # would silently receive no gradient, so the model trains as if `lm_head` were frozen. Fail loudly rather + # than train a silently-frozen head. + output_embeddings = model.get_output_embeddings() + if isinstance(output_embeddings, BaseTunerLayer): + raise ValueError( + "`use_liger_kernel=True` is incompatible with applying a PEFT adapter to `lm_head`. The Liger " + "fused GRPO loss reads `lm_head.weight` directly, so the adapter on the head is ignored and never " + "trained. Either remove `'lm_head'` from your `target_modules`, or set `use_liger_kernel=False`." + ) + self.mask_truncated_completions = args.mask_truncated_completions + self.top_entropy_quantile = args.top_entropy_quantile + if self.use_liger_kernel and self.top_entropy_quantile < 1.0: + raise NotImplementedError( + "Liger Kernels don't currently support masking token positions based on entropy." + ) + if self.use_liger_kernel and self.importance_sampling_level not in ("token", "sequence"): + raise ValueError( + f"Unknown importance sampling level: {self.importance_sampling_level}. " + "Possible values are 'token' and 'sequence'." + ) + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + + if train_dataset is None: + raise ValueError("`train_dataset` is required") + elif ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in GRPOTrainer. Please use a standard dataset instead." + ) + + if args.loss_type == "luspo" and args.importance_sampling_level != "sequence": + logger.warning( + "When using `'luspo'` loss, `importance_sampling_level` should be set to `'sequence'` to mirror the " + "paper's setup." + ) + + if args.loss_type == "vespo" and args.importance_sampling_level != "token": + logger.warning( + "VESPO computes sequence-level importance weights internally. `importance_sampling_level` should be " + "set to `'token'` (the default)." + ) + + if args.importance_sampling_level == "sequence" and args.loss_type in ["bnpo", "dr_grpo", "dapo", "cispo"]: + logger.warning( + f"When using `importance_sampling_level='sequence'`, the `'{args.loss_type}'` loss sums per-token " + "contributions, which effectively weights each sequence by its completion length instead of " + "optimizing the per-sequence objective. To reproduce the GSPO paper's setup, set `loss_type='grpo'` " + "(see https://huggingface.co/docs/trl/main/en/paper_index#group-sequence-policy-optimization]." + ) + + if self.loss_type == "vespo" and self.use_vllm and self.vllm_importance_sampling_correction: + if self.vllm_importance_sampling_mode not in ["token_truncate", "token_mask"]: + raise ValueError( + f"VESPO loss requires `vllm_importance_sampling_mode` to be either 'token_truncate' or " + f"'token_mask'. Got: {self.vllm_importance_sampling_mode}." + ) + + # Multi-step + self.num_iterations = args.num_iterations # = 𝜇 in the GRPO paper + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + # Tracks the number of iterations [forward + backward passes], including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in GRPO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + # In Trainer, `training_step` scales the loss by `gradient_accumulation_steps` only if `compute_loss_func` + # is None. For DAPO, loss scaling instead depends on the total number of completions tokens across the + # global accumulated batch. To control scaling ourselves, we must disable Trainer’s built-in scaling. The + # simplest [though a bit hacky] way is to set `compute_loss_func` to any non-None value, which bypasses + # that behavior without rewriting `training_step`. + compute_loss_func="non-None value to disable scaling", + ) + + # Reference model + self.beta = args.beta + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + self.ref_model = create_model_from_path(get_config_model_id(self.model.config), **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Cast LM Head To FP32 + if args.cast_lm_head_to_fp32: + + def _cast_lm_head_to_fp32(target_model: PreTrainedModel): + """Cast lm_head to fp32 while preserving embedding output dtype if tied.""" + + def cast_inputs_to_fp32(module, inputs): + # Preserve other positional args and kwargs untouched + if not inputs: + return inputs + return (inputs[0].to(torch.float32),) + inputs[1:] + + original_dtype_local = target_model.lm_head.weight.dtype + target_model.lm_head = target_model.lm_head.float() + target_model.lm_head.register_forward_pre_hook(cast_inputs_to_fp32) + + if target_model.config.tie_word_embeddings: + + def cast_outputs_to_original_dtype(module, args, output): + return output.to(original_dtype_local) + + # Only cast activations; weights are now fp32 [intentional for numerical stability of logits] + target_model.model.embed_tokens.register_forward_hook(cast_outputs_to_original_dtype) + + _cast_lm_head_to_fp32(model) + if self.ref_model is not None: + _cast_lm_head_to_fp32(self.ref_model) + + # Liger loss + if self.use_liger_kernel: + if not is_liger_kernel_available(): + raise ImportError( + "Liger is required to use `use_liger_kernel` as the GRPO loss. Run `pip install liger-kernel`." + ) + # redirect the model.module forward to the model forward to ensure pre-forward hooks are called + self._forward_redirection = _ForwardRedirection() + + self.liger_grpo_loss = LigerFusedLinearGRPOLoss( + beta=self.beta, + epsilon_low=self.epsilon_low, + epsilon_high=self.epsilon_high, + temperature=self.temperature, + use_ref_model=self.beta != 0.0, + loss_type=self.loss_type, + max_completion_length=self.max_completion_length, + importance_sampling_level=self.importance_sampling_level, + delta=args.delta, + use_bias_correction_kl=args.use_bias_correction_kl, + sapo_temperature_pos=args.sapo_temperature_pos, + sapo_temperature_neg=args.sapo_temperature_neg, + vespo_k_pos=args.vespo_k_pos, + vespo_lambda_pos=args.vespo_lambda_pos, + vespo_k_neg=args.vespo_k_neg, + vespo_lambda_neg=args.vespo_lambda_neg, + ) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self._current_train_step_time = 0.0 + self.log_completions = args.log_completions + self.log_unique_prompts = args.log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # Keep logs sized to the generation batch to record only outputs from the latest model update. + self._logs = { + "images": deque(maxlen=args.generation_batch_size), + "prompt": deque(maxlen=args.generation_batch_size), + "completion": deque(maxlen=args.generation_batch_size), + "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + "advantages": deque(maxlen=args.generation_batch_size), + "extra": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + } + # Buffers for user-logged data from reward functions, flushed after gathering + self._pending_extra_logs = defaultdict(list) + self._pending_metrics = defaultdict(list) + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + self.vllm_generation = VLLMGeneration( + model=self.model, + accelerator=self.accelerator, + processing_class=self.processing_class, + mode=args.vllm_mode, + structured_outputs_regex=args.vllm_structured_outputs_regex, + server_base_url=args.vllm_server_base_url, + server_host=args.vllm_server_host, + server_port=args.vllm_server_port, + group_port=args.vllm_group_port, + server_timeout=args.vllm_server_timeout, + tensor_parallel_size=args.vllm_tensor_parallel_size, + gpu_memory_utilization=args.vllm_gpu_memory_utilization, + max_model_length=args.vllm_max_model_length, + max_num_seqs=args.per_device_train_batch_size + * args.vllm_tensor_parallel_size + * args.steps_per_generation, + enable_sleep_mode=args.vllm_enable_sleep_mode, + model_impl=args.vllm_model_impl, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + min_p=self.min_p, + max_completion_length=self.max_completion_length, + logprobs=0, + generation_kwargs=args.generation_kwargs, + ) + self._last_loaded_step = -1 + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": self._tokenizer.pad_token_id, + "bos_token_id": self._tokenizer.bos_token_id, + "eos_token_id": self._tokenizer.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs, disable_compile=True) + # Keep training-specific generation kwargs to overwrite model's original generation config + self.generation_kwargs = generation_kwargs + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + self._dist = DistributedBackend(self.accelerator) + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + if self.beta == 0.0: + raise ValueError( + "You passed `sync_ref_model=True` while `beta=0.0`, which means the reference model is not used " + "during training. Consequently, GRPOTrainer does not create a `ref_model` instance, and there is " + "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." + ) + if is_peft_model(model): + raise NotImplementedError( + "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " + "With PEFT, GRPOTrainer does not keep a separate reference model in memory; instead, it recovers " + "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " + "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " + "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " + "feature request at https://github.com/huggingface/trl/issues." + ) + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + if self.accelerator.is_main_process and self.log_completions: + os.makedirs(os.path.join(self.args.output_dir, "completions"), exist_ok=True) + if self.args.log_completions_hub_repo is not None: + repo_id = self.args.log_completions_hub_repo + create_repo(repo_id, private=self.args.hub_private_repo, repo_type="dataset", exist_ok=True) + template_path = pkg_resources.files("trl").joinpath("templates/completions_dataset_card.md") + card_data = DatasetCardData( + pretty_name="TRL Completion logs", + tags=["trl", "trl-logs", "completions"], + ) + card = DatasetCard.from_template( + card_data=card_data, + template_path=str(template_path), + repo_id=repo_id, + hub_model_id=self.args.hub_model_id, + ) + card.push_to_hub(repo_id) + self.commit_scheduler = CommitScheduler( + repo_id=repo_id, + repo_type="dataset", + folder_path=f"{self.args.output_dir}/completions", + every=2, # minutes + allow_patterns=["*.parquet"], + ) + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't + # work. Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt", "image", "images"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. + def get_train_dataloader(self): + return self._get_dataloader( + dataset=self.train_dataset, + description="Training", + batch_size=self._train_batch_size * self.args.steps_per_generation, # < this is the change + sampler_fn=self._get_train_sampler, + is_training=True, + ) + + def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. Each row shows the per-step batch + # returned by `_prepare_inputs`; rows within a `steps_per_generation` block are slices of the same + # generated batch. When `num_iterations > 1`, that block is reused for multiple optimization passes + # before regenerating. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations_eval, + seed=self.args.seed, + ) + + @profiling_decorator + def _get_last_hidden_state( + self, + unwrapped_model, + input_ids, + attention_mask, + logits_to_keep, + pixel_values=None, + image_grid_thw=None, + pixel_attention_mask=None, + spatial_shapes=None, + image_sizes=None, + image_position_ids=None, + ): + if is_peft_model(unwrapped_model): + unwrapped_model = unwrapped_model.base_model.model + + # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) + model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask} + + # For Qwen models: + if image_grid_thw is not None and pixel_values is not None: + model_inputs["image_grid_thw"] = image_grid_thw + # For Gemma, SmolVLM2, LLaVa-Next etc.: + if pixel_values is not None: + model_inputs["pixel_values"] = pixel_values + # For SmolVLM2 + if pixel_attention_mask is not None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask + # For LFM2-VL + if spatial_shapes is not None: + model_inputs["spatial_shapes"] = spatial_shapes + # For LLaVa-Next + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes + if image_position_ids is not None: + model_inputs["image_position_ids"] = image_position_ids + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + # `base_model` gives the backbone model (skipping `lm_head`) — text decoder for LMs, multimodal wrapper for + # VLMs (so vision-token injection runs before the text decoder). `get_decoder()` won't do: on VLMs it + # returns just the text stack and feeds image-placeholder IDs through it. + # Pre-5.0 transformers VLMs set `base_model_prefix = ""` so `base_model is self` (re-runs `lm_head`). + # Fall back to `.model` there. + if self._is_vlm and Version(transformers.__version__) < Version("5.0.0"): + backbone = unwrapped_model.model + else: + backbone = unwrapped_model.base_model + last_hidden_state = backbone(**model_inputs).last_hidden_state + # Exclude the last value: it corresponds to the next token pred + last_hidden_state = last_hidden_state[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + last_hidden_state = last_hidden_state[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + return last_hidden_state + + def get_high_entropy_mask(self, entropies: torch.Tensor, mask: torch.Tensor, threshold: float) -> torch.Tensor: + """ + Returns a binary mask identifying tokens whose entropy exceeds a given quantile threshold. + + Args: + entropies (`torch.Tensor`): + Tensor of shape (batch_size, seq_len) with per-token entropy values. + mask (`torch.Tensor`): + Binary mask of the same shape as `entropies`, where `1` indicates valid tokens and `0` padding. + threshold (`float`): + Quantile threshold between `0.0` and `1.0` to select high-entropy tokens. + + Returns: + `torch.Tensor`: + Boolean mask of shape (batch_size, seq_len), where `True` indicates tokens with entropy >= threshold + and `False` otherwise. + """ + local = entropies[mask.bool()].float() + + # Use a negative pad_value as a sentinel because entropy values are always >= 0. + # This guarantees that the sentinel cannot collide with any real entropy value. + pad_value = -1e9 + + # Pad across processes so that every rank has the same tensor length + padded = self.accelerator.pad_across_processes(local, dim=0, pad_index=pad_value) + gathered = self.accelerator.gather(padded) + + # Drop sentinel values (safe because no entropy can be negative) + gathered = gathered[gathered != pad_value] + + if gathered.numel() == 0: + return torch.zeros_like(entropies, dtype=torch.bool) + + entropy_threshold = torch.quantile(gathered, threshold) + masked_entropies = entropies * mask.float() + entropy_mask = masked_entropies >= entropy_threshold + return entropy_mask & mask.bool() # ensure padding tokens are always masked out + + def _get_per_token_logps_and_entropies( + self, + model, + input_ids, + attention_mask, + logits_to_keep, + batch_size = None, + compute_entropy = False, + compute_efficient = False, + *args, + **kwargs, + ): + # All Unsloth code here in this function is licensed under AGPL3 + # if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': + # return None, None # logps, entropies Unsloth efficient GRPO + if compute_efficient: + return None, None + else: + if not hasattr(self, "_autocast_dtype"): + self._autocast_dtype = ( + torch.float16 + if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16" + else torch.bfloat16 + ) + if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": + self._autocast_dtype = torch.float16 + + pixel_values, image_grid_thw = ( + kwargs.get("pixel_values", None), + kwargs.get("image_grid_thw", None), + ) + pixel_attention_mask, image_sizes = ( + kwargs.get("pixel_attention_mask", None), + kwargs.get("image_sizes", None), + ) + num_images = kwargs.get("num_images", None) + # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models + token_type_ids = kwargs.get("token_type_ids", None) + mm_token_type_ids = kwargs.get("mm_token_type_ids", None) + if mm_token_type_ids is not None or image_grid_thw is not None: + mm_token_type_ids = _unsloth_fix_mm_token_type_ids( + self.processing_class, input_ids, mm_token_type_ids + ) + + unwrapped_model = self.accelerator.unwrap_model(model, keep_fp32_wrapper = False) + + lm_head = self.model.get_output_embeddings().weight + + dtype_bytes = 16 if self._autocast_dtype in [torch.float16, torch.bfloat16] else 32 + total_rows = input_ids.shape[0] + seq_len = input_ids.shape[1] + hidden_dim = lm_head.shape[1] + vocab_dim = lm_head.shape[0] + + if self.args.unsloth_grpo_mini_batch is None: + B, multiplier = autotune_batch_and_chunks( + total_rows, + seq_len, + hidden_dim, + vocab_dim, + dtype_bytes, + self.args.unsloth_logit_chunk_multiplier, + ) + B = total_rows // B + else: + B = self.args.unsloth_grpo_mini_batch + + if self.args.unsloth_logit_chunk_multiplier is None: + multiplier = max(4, seq_len // 4096) + else: + multiplier = self.args.unsloth_logit_chunk_multiplier + + all_logprobs_list = [] + if pixel_values is None: + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt( + input_ids, logits_to_keep, self.processing_class.pad_token_id + ) + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() + input_ids = left_pack_padding(input_ids, self.processing_class.pad_token_id) + attention_mask = input_ids != self.processing_class.pad_token_id + attention_mask = attention_mask.to(attention_mask.dtype) + else: + max_left_pad = 0 + + def slice_sample_axis(value, start, end): + if value is None: + return None + return value[start:end] + + import math + + total_samples = input_ids.shape[0] + batch_size = math.ceil(total_samples / B) + if isinstance(num_images, torch.Tensor): + num_images = num_images.detach().cpu().reshape(-1).tolist() + if image_grid_thw is not None and pixel_values is not None and num_images is not None: + rows_per_image = image_grid_thw.prod(dim = -1) + rows_per_sample = torch.split(rows_per_image, num_images) + rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) + # why: cum_rows is indexed via .item() inside the per-chunk loop; + # keeping it on CPU avoids per-iteration GPU->CPU sync. + cum_rows = torch.cat( + [ + torch.tensor([0], device = rows_per_sample.device), + rows_per_sample.cumsum(0), + ] + ).cpu() + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + else: + cum_rows = None + cum_imgs = None + + def _first_dim_len(value): + if value is None: + return None + if hasattr(value, "shape"): + return value.shape[0] + try: + return len(value) + except TypeError: + return None + + total_images = sum(num_images) if num_images is not None else None + _image_sizes_n = _first_dim_len(image_sizes) + + input_ids_chunks = [] + attention_mask_chunks = [] + pixel_values_chunks = [] + image_grid_thw_chunks = [] + pixel_attention_mask_chunks = [] + image_sizes_chunks = [] + token_type_ids_chunks = [] + mm_token_type_ids_chunks = [] + + current_pixel_idx = 0 + # TRL 0.23.0 batching logic + for start in range(0, total_samples, batch_size): + end = min(start + batch_size, total_samples) + + input_ids_chunks.append(input_ids[start:end]) + attention_mask_chunks.append(attention_mask[start:end]) + token_type_ids_chunks.append(slice_sample_axis(token_type_ids, start, end)) + mm_token_type_ids_chunks.append(slice_sample_axis(mm_token_type_ids, start, end)) + + if image_grid_thw is not None and pixel_values is not None: + if num_images is None: + grid_slice = image_grid_thw[start:end] + batch_pixel_count = grid_slice.prod(dim = -1).sum().item() + start_pixel_idx = current_pixel_idx + end_pixel_idx = current_pixel_idx + batch_pixel_count + current_pixel_idx = end_pixel_idx + img_start = img_end = None + else: + start_pixel_idx = cum_rows[start].item() + end_pixel_idx = cum_rows[end].item() + img_start = cum_imgs[start].item() + img_end = cum_imgs[end].item() + grid_slice = image_grid_thw[img_start:img_end] + image_grid_thw_chunks.append(grid_slice) + + pixel_values_chunks.append(pixel_values[start_pixel_idx:end_pixel_idx]) + + if image_sizes is None: + image_sizes_chunks.append(None) + elif ( + num_images is not None + and _image_sizes_n == total_images + and img_start is not None + ): + image_sizes_chunks.append(image_sizes[img_start:img_end]) + else: + image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) + + if pixel_attention_mask is None: + pixel_attention_mask_chunks.append(None) + elif ( + num_images is not None + and img_start is not None + and pixel_attention_mask.shape[0] == image_grid_thw.shape[0] + ): + pixel_attention_mask_chunks.append(pixel_attention_mask[img_start:img_end]) + elif ( + pixel_attention_mask.shape[0] == pixel_values.shape[0] + and pixel_attention_mask.shape[0] != input_ids.shape[0] + ): + pixel_attention_mask_chunks.append( + pixel_attention_mask[start_pixel_idx:end_pixel_idx] + ) + else: + pixel_attention_mask_chunks.append(pixel_attention_mask[start:end]) + + else: + pixel_values_chunks.append(None) + image_grid_thw_chunks.append(None) + pixel_attention_mask_chunks.append(None) + image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) + + temperature = self.temperature + logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) + logit_scale_multiply = getattr(model.config, "logit_scale", 0) + if logit_scale_multiply is None: + logit_scale_multiply = 0 + logit_scale_divide = getattr(model.config, "logits_scaling", 0) + if logit_scale_divide is None: + logit_scale_divide = 0 + + zipped_inputs = zip( + input_ids_chunks, + attention_mask_chunks, + pixel_values_chunks, + image_grid_thw_chunks, + pixel_attention_mask_chunks, + image_sizes_chunks, + token_type_ids_chunks, + mm_token_type_ids_chunks, + ) + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + + with _get_inference_mode_context_manager(model): + for ( + input_ids_chunk, + attention_mask_chunk, + pixel_values_chunk, + image_grid_thw_chunk, + pixel_attention_mask_chunk, + image_sizes_chunk, + token_type_ids_chunk, + mm_token_type_ids_chunk, + ) in zipped_inputs: + _extra_vision_kwargs = {} + if token_type_ids_chunk is not None: + _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk + if mm_token_type_ids_chunk is not None: + _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk + with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): + if pixel_values is None: + logits_chunk = unwrapped_model( + input_ids = input_ids_chunk, + attention_mask = attention_mask_chunk, + pixel_values = pixel_values_chunk, + image_grid_thw = image_grid_thw_chunk, + pixel_attention_mask = pixel_attention_mask_chunk, + image_sizes = image_sizes_chunk, + **_extra_vision_kwargs, + ).logits + + completion_input_ids_chunk = input_ids_chunk[ + :, -(logits_to_keep + max_left_pad) : + ] + logits_chunk = logits_chunk[ + :, -(logits_to_keep + max_left_pad + 1) :, : + ] + logits_chunk = logits_chunk[:, :-1, :] + logprobs_chunk = chunked_hidden_states_selective_log_softmax( + logits_chunk, + lm_head, + completion_input_ids_chunk, + chunks = input_ids_chunk.shape[0] * multiplier, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + logit_softcapping = logit_softcapping, + temperature = temperature, + ) + else: + # Essentially, for VLMs we do not go via the optimized path in models/, + # so we don't encounter the Flash Attn left-padding issue. + logits_chunk = unwrapped_model( + input_ids = input_ids_chunk, + attention_mask = attention_mask_chunk, + pixel_values = pixel_values_chunk, + image_grid_thw = image_grid_thw_chunk, + pixel_attention_mask = pixel_attention_mask_chunk, + image_sizes = image_sizes_chunk, + logits_to_keep = logits_to_keep + 1, + **_extra_vision_kwargs, + ).logits + + logits_chunk = logits_chunk[:, :-1, :] + completion_input_ids_chunk = input_ids_chunk[:, -logits_to_keep:] + # Guard: check if model returned hidden states or logits + if logits_chunk.shape[-1] == lm_head.shape[1]: + logprobs_chunk = chunked_hidden_states_selective_log_softmax( + logits_chunk, + lm_head, + completion_input_ids_chunk, + chunks = input_ids_chunk.shape[0] * multiplier, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + logit_softcapping = logit_softcapping, + temperature = temperature, + ) + else: + # Model returned logits directly - scaling/softcapping already applied by model forward + logprobs_chunk = chunked_selective_log_softmax( + logits_chunk, + completion_input_ids_chunk, + temperature, + ) + # This is needed to avoid race conditions with GPT OSS offload_embbed=True + # However, it seems that this line does not slow down or disrupt models. + device_synchronize() + all_logprobs_list.append(logprobs_chunk) + logprobs = torch.cat(all_logprobs_list, dim = 0) + entropies = None + + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" + + return logprobs.detach(), 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 + # logits = logits[:, -logits_to_keep:] + # return logits + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + # logits = logits / self.temperature + # logps = selective_log_softmax(logits, input_ids) + + # row_indices, col_indices = torch.where(logps < -20) + + # # Method 1: Check if tensors have elements + # if len(row_indices) > 0 and len(col_indices) > 0: + # breakpoint() # Breakpoint triggered here + # print("Found high values!") + # return logps # compute logprobs for the input tokens + + def training_step(self, model, inputs, num_items_in_batch): + time_before = time.perf_counter() + output = super().training_step(model, inputs, num_items_in_batch) + self._step += 1 + time_after = time.perf_counter() + self._current_train_step_time += time_after - time_before + if self._step % self.current_gradient_accumulation_steps == 0: + self._metrics["train"]["step_time"].append(self._current_train_step_time) + self._current_train_step_time = 0.0 + return output + + @profiling_decorator + def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = split_pixel_values_by_grid(generation_batch) + + try: generation_batch = shuffle_sequence_dict(generation_batch) + + except: pass + generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) + self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + return inputs + + def _log_completion_extra(self, column: str, values: list): + """ + Log extra columns to the completions table. Called from reward functions via the `log_extra` kwarg. + + Args: + column (`str`): + Name of the column to add. + values (`list`): + Values for the column, one per sample in the batch. + """ + self._pending_extra_logs[column].extend(values) + + def _log_metric(self, name: str, value: float): + """ + Log a scalar metric from a reward function. Called via the `log_metric` kwarg. Values are averaged over each + logging step and reported alongside built-in metrics like `kl` and `entropy`. + + Args: + name (`str`): + Name of the metric. + value (`float`): + Scalar value for this batch. + """ + self._pending_metrics[name].append(value) + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + # This allows for dynamic reward shaping based on training progress. + reward_kwargs["trainer_state"] = self.state + + # Allow reward functions to log extra columns to the completions table. + reward_kwargs["log_extra"] = self._log_completion_extra + + # Allow reward functions to log additional scalar metrics. + reward_kwargs["log_metric"] = self._log_metric + + async_funcs_info = [] # async custom functions for asyncio.gather + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) + ): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + with profiling_context(self, reward_func_name): + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] + texts = [ + apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] + for x in messages + ] + else: + texts = [p + c for p, c in zip(prompts, completions, strict=True)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + elif inspect.iscoroutinefunction(reward_func): # Separate async reward funcs to run them in parallel later + async_funcs_info.append((i, reward_func, reward_func_name)) + else: + # Run synchronous reward function + with profiling_context(self, reward_func_name): + if self.environments is not None: + reward_kwargs["environments"] = self.environments + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # Execute async custom functions in parallel using asyncio.gather + if async_funcs_info: + + async def _invoke_async(index, func, func_name): + with profiling_context(self, func_name): + output = await func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + output = [r if r is not None else torch.nan for r in output] + return index, output + + async def _run_async_funcs(): + coros = [_invoke_async(i, func, func_name) for (i, func, func_name) in async_funcs_info] + return await asyncio.gather(*coros) + + async_results = asyncio.run_coroutine_threadsafe(_run_async_funcs(), self.async_loop).result() + for idx, output_reward_func in async_results: + rewards_per_func[:, idx] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = { + key: value[nan_row_idx] + for key, value in reward_kwargs.items() + if key not in ("trainer_state", "log_extra", "log_metric") + } + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + logger.warning( + f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _tokenize_prompts(self, prompts: list): + """Tokenize prompts and extract images/multimodal fields for generation.""" + if is_conversational({"prompt": prompts[0]}): + # Normalize string content to content blocks for VLM processors that don't handle plain strings. + if self._is_vlm: + prompts = [prepare_multimodal_messages(prompt) for prompt in prompts] + + # Extract images from messages for VLM support + images = [] + has_images = False + for prompt in prompts: + prompt_images = [] + for message in prompt: + if isinstance(message["content"], list): + for part in message["content"]: + if part["type"] == "image": + prompt_images.append(part["image"]) + has_images = True + images.append(prompt_images if prompt_images else None) + images = images if has_images else None + + # Workaround for a bug in transformers 5.3.0 where some processors (e.g. Qwen2.5-VL) crash on + # batched unpadded input (transformers#44514). + # Fixed in transformers 5.4.0 (transformers#44563). + needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") + tokenized = self.processing_class.apply_chat_template( + conversation=prompts, + tools=self.tools or None, # `or None`: Llama bug: it renders tool boilerplate for tools=[] + chat_template=self.chat_template, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **({"padding": True} if needs_padding_workaround else {}), + **self.chat_template_kwargs, + ) + if needs_padding_workaround: + # Unpad input_ids: remove padding tokens using attention_mask to get per-sequence lists + prompt_ids = [ + [tok for tok, m in zip(ids, mask, strict=True) if m] + for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) + ] + else: + prompt_ids = tokenized["input_ids"] + # For VLMs, the processor returns extra multimodal fields (pixel_values, image_grid_thw, etc.) + multimodal_fields = {k: v for k, v in tokenized.items() if k not in ("input_ids", "attention_mask")} + else: + prompt_ids = self.processing_class(text=prompts)["input_ids"] + images = None + multimodal_fields = {} + return prompt_ids, images, multimodal_fields + + def _generate_single_turn(self, prompt_ids, images, multimodal_fields): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + # Sync weights if training step changed + if self.state.global_step != self._last_loaded_step: + if not getattr(getattr(self.vllm_generation, 'llm', None), 'shared_weights', False): + with profiling_context(self, 'sync_weights'): + self.vllm_generation.sync_weights() + self._last_loaded_step = self.state.global_step + + # Generate using vLLM with raw token IDs + num_generations = self.num_generations if mode == "train" else self.num_generations_eval + _, completion_ids, logprobs, _ = self.vllm_generation.generate( + prompts=prompt_ids, + images=images, + num_generations=num_generations, + profiler=profiling_context(self, "vLLM.generate"), + ) + # vLLM returns per-token top-k logprobs; keep only the top-1 (sampled token) logprob + logprobs = [[lp[0] for lp in seq] for seq in logprobs] + + elif self.use_transformers_continuous_batching: + with ( + profiling_context(self, "transformers.generate_batch"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + self._dist.summon_full_params(self.model_wrapped, recurse=False), + ): + # Cast to the appropriate dtype based on training configuration + if self.args.bf16: + unwrapped_model.to(torch.bfloat16) + elif self.args.fp16: + unwrapped_model.to(torch.float16) + if self.args.cast_lm_head_to_fp32: + unwrapped_model.lm_head.to(torch.float32) + all_outputs = unwrapped_model.generate_batch( + prompt_ids, + generation_config=self.generation_config, + continuous_batching_config=self.continuous_batching_config, + progress_bar=False, + ) + unwrapped_model.train() + completion_ids = [output.generated_tokens for output in all_outputs.values()] + logprobs = None + + else: + # Regular generation path: left-pad token IDs into tensors + prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] + padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left") + attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left") + generate_inputs = {"input_ids": padded_ids, "attention_mask": attention_mask} + # For VLMs, include multimodal fields as tensors (pixel_values, image_grid_thw, etc.) + for k, v in multimodal_fields.items(): + if isinstance(v, torch.Tensor): + generate_inputs[k] = v + elif isinstance(v, list) and v and isinstance(v[0], list): + # Per-token field (e.g., token_type_ids): left-pad like input_ids + generate_inputs[k] = pad([torch.tensor(x) for x in v], padding_value=0, padding_side="left") + else: + generate_inputs[k] = torch.tensor(np.array(v)) + generate_inputs = super()._prepare_inputs(generate_inputs) + if "mm_token_type_ids" in generate_inputs or "image_grid_thw" in generate_inputs: + mm_token_type_ids = _unsloth_fix_mm_token_type_ids( + self.processing_class, + generate_inputs["input_ids"], + generate_inputs.get("mm_token_type_ids", None), + ) + if mm_token_type_ids is not None: + generate_inputs["mm_token_type_ids"] = mm_token_type_ids + + with ( + profiling_context(self, "transformers.generate"), + unwrap_model_for_generation( + self.model_wrapped, + self.accelerator, + gather_deepspeed3_params=self.args.ds3_gather_for_generation, + generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 + ) as unwrapped_model, + torch.no_grad(), + self._dist.summon_full_params(self.model_wrapped, recurse=False), + ): + prompt_completion_ids = unwrapped_model.generate( + **generate_inputs, generation_config=self.generation_config + ) + # Compute prompt length and extract completion ids + prompt_length = generate_inputs["input_ids"].size(1) + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self._tokenizer.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + completion_ids = [ + c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) + ] + logprobs = None # not used in this case + + return completion_ids, logprobs + + def _get_tool_suffix_ids(self, tool_messages): + """Get token IDs for tool result formatting by using a minimal dummy conversation.""" + # Use the real tool name instead of a dummy: some templates (e.g. GPT-OSS) derive the tool response + # header from the assistant's tool call name. + dummy_tool_calls = [{"type": "function", "function": {"name": tool_messages[0]["name"], "arguments": {}}}] + dummy_messages = [ + {"role": "user", "content": "dummy"}, + { + "role": "assistant", + # "content" is required here because VLM processors crash on tokenize=True without it + # (KeyError in processing_utils.py). See huggingface/transformers#45290. + "content": "", + "tool_calls": dummy_tool_calls, + }, + ] + if self._is_vlm: + dummy_messages = prepare_multimodal_messages(dummy_messages) + tool_messages = prepare_multimodal_messages(tool_messages) + + prefix_ids = self.processing_class.apply_chat_template( + dummy_messages, + add_generation_prompt=False, + tokenize=True, + chat_template=self.chat_template, + return_dict=False, + **self.chat_template_kwargs, + ) + full_ids = self.processing_class.apply_chat_template( + dummy_messages + tool_messages, + add_generation_prompt=True, + tokenize=True, + chat_template=self.chat_template, + return_dict=False, + **self.chat_template_kwargs, + ) + # VLM processors return batched output (list of lists), unbatch for single conversation + if self._is_vlm: + prefix_ids = prefix_ids[0] + full_ids = full_ids[0] + + # Some chat templates (notably Qwen3/Qwen3.5) render "...<|im_end|>\n" after an assistant/tool block. + # When we compute `suffix_ids` by slicing `full_ids`, we must align the slicing boundary to + # EOS (not EOS + newline). Templates that don't use EOS as end-of-turn (e.g. Gemma uses + # ) skip this trimming. + eos_positions = [i for i, tok_id in enumerate(prefix_ids) if tok_id == self._tokenizer.eos_token_id] + if eos_positions: + prefix_ids = prefix_ids[: eos_positions[-1] + 1] + + if full_ids[: len(prefix_ids)] != prefix_ids: + raise ValueError("Unexpected tokenization: the EOS-trimmed prefix IDs are not a prefix of the full IDs.") + return full_ids[len(prefix_ids) :] + + def _tool_call_loop(self, prompts, prompt_ids, completion_ids, completions, logprobs, images, multimodal_fields): + # Tool execution loop: execute tools, then regenerate completions with tool results appended to the prompt + tool_calls = [completion[0].get("tool_calls") for completion in completions] + idxs_with_tool = [idx for idx, tool_call in enumerate(tool_calls) if tool_call] + tool_calls = [tool_calls[idx] for idx in idxs_with_tool] + tool_mask = [[1] * len(ids) for ids in completion_ids] # 0 for tool result tokens, 1 elsewhere + # Collect images from multimodal tool responses for the forward pass + tool_images = [[] for _ in completion_ids] + tool_call_count = 0 + tool_failure_count = 0 + iteration_num = 0 + + while idxs_with_tool and iteration_num < self.max_tool_calling_iterations: + prompt_completion_tools = [prompts[i] for i in idxs_with_tool] # select only prompts that need tool calls + # Snapshot state so we can rollback tool results that would exceed max_completion_length + completions_len_before = [len(completions[i]) for i in idxs_with_tool] + tool_images_len_before = [len(tool_images[i]) for i in idxs_with_tool] + prompts_len_before = [len(prompts[i]) for i in idxs_with_tool] + + # Call the tools, and build the new prompt for generation + for idx in range(len(idxs_with_tool)): + idx_with_tool = idxs_with_tool[idx] + tool_call_list = tool_calls[idx] + prompt_completion_tool = prompt_completion_tools[idx] + sync_tool_dict = self._sync_tool_dicts[idx_with_tool] + async_tool_dict = self._async_tool_dicts[idx_with_tool] + # Append the last assistant message (which triggered tool_calls) to the prompt + prompt_completion_tool.append(completions[idx_with_tool][-1]) + async_coros = [] + tool_call_results = [] + for tool_call in tool_call_list: + tool_call_count += 1 + if tool_call["type"] == "function": + function = tool_call["function"] + name = function["name"] + try: + if name in sync_tool_dict: + tool_call_results.append((name, sync_tool_dict[name](**function["arguments"]))) + elif name in async_tool_dict: + async_coros.append((name, async_tool_dict[name](**function["arguments"]))) + else: + raise ValueError(f"Tool {name} not found.") + except Exception as e: + tool_failure_count += 1 + result = {"error": str(e)} + tool_call_results.append((name, result)) + else: + tool_failure_count += 1 + name = tool_call.get("name", "unknown") + tool_call_results.append((name, {"error": f"Unsupported tool call type: {tool_call['type']}"})) + + if async_coros: + + async def _run_async_tools(async_coros): + coros = [coro for _, coro in async_coros] + results = await asyncio.gather(*coros, return_exceptions=True) + return [(name, result) for (name, _), result in zip(async_coros, results, strict=False)] + + async_results = asyncio.run_coroutine_threadsafe( + _run_async_tools(async_coros), self.async_loop + ).result() + + for name, result in async_results: + if isinstance(result, Exception): + tool_failure_count += 1 + tool_call_results.append((name, {"error": str(result)})) + else: + tool_call_results.append((name, result)) + + for name, result in tool_call_results: + # Support multimodal tool responses: if the tool returns a list of content blocks + # (e.g., [{"type": "image", "image": ...}, {"type": "text", "text": "..."}]), + # pass them through directly so _tokenize_prompts can extract images for VLMs. + content = result if isinstance(result, list) else str(result) + tool_message = {"role": "tool", "name": name, "content": content} + # Collect images from multimodal tool responses + if isinstance(content, list): + for part in content: + if isinstance(part, dict) and part.get("type") == "image": + tool_images[idx_with_tool].append(part["image"]) + prompt_completion_tool.append(tool_message) + completions[idx_with_tool].append(tool_message) + + # Build token IDs by concatenation: prompt + completion + tool_suffix. + prompt_completion_tool_ids = [] + for idx in range(len(idxs_with_tool)): + idx_with_tool = idxs_with_tool[idx] + # Extract trailing tool messages from completions + tool_messages = [] + for message in reversed(completions[idx_with_tool]): + if message["role"] == "tool": + tool_messages.insert(0, message) + else: + break + suffix_ids = self._get_tool_suffix_ids(tool_messages) + prompt_completion_tool_ids.append( + prompt_ids[idx_with_tool] + completion_ids[idx_with_tool] + suffix_ids + ) + + # Drop tool results whose addition would push the sequence past max_completion_length (the completion + # budget) or past the backend context ceiling (vLLM and transformers will error out on inputs longer than + # the model's max length). The sample exits the loop with its completion as-is, and the tool + # messages/images appended this iteration are rolled back so completions and tool_images stay consistent + # with completion_ids. + if self.use_vllm and self.vllm_mode == "colocate": + max_model_len = self.vllm_generation.llm.llm_engine.model_config.max_model_len + else: + config = self.model.config.text_config if self._is_vlm else self.model.config + max_model_len = config.max_position_embeddings + overlong = [ + len(pct) - len(prompt_ids[i]) > self.max_completion_length or len(pct) >= max_model_len + for i, pct in zip(idxs_with_tool, prompt_completion_tool_ids, strict=True) + ] + for idx in range(len(idxs_with_tool)): + if overlong[idx]: + idx_with_tool = idxs_with_tool[idx] + del completions[idx_with_tool][completions_len_before[idx] :] + del tool_images[idx_with_tool][tool_images_len_before[idx] :] + del prompts[idx_with_tool][prompts_len_before[idx] :] + # Keep only non-overlong items for further processing + idxs_with_tool = [idx for idx, o in zip(idxs_with_tool, overlong, strict=True) if not o] + prompt_completion_tool_ids = [ + pct for pct, o in zip(prompt_completion_tool_ids, overlong, strict=True) if not o + ] + if not idxs_with_tool: + break # all overlong, exit tool loop + + # Filter images and multimodal fields to match the current subset (index into full batch). + # Merge tool response images so the model can see visual feedback during generation. + merged_images = images + if any(imgs for imgs in tool_images): + if merged_images is None: + merged_images = [imgs if imgs else None for imgs in tool_images] + else: + merged_images = [ + (existing or []) + new for existing, new in zip(merged_images, tool_images, strict=True) + ] + loop_images = [merged_images[i] for i in idxs_with_tool] if merged_images else None + if multimodal_fields: + loop_multimodal_fields = {} + for k, v in multimodal_fields.items(): + selected = [v[i] for i in idxs_with_tool] + # Per-token fields (e.g. token_type_ids) need zero-padding to match extended prompt length + if isinstance(selected[0], list): + selected = [ + s + [0] * (len(pct) - len(s)) + for s, pct in zip(selected, prompt_completion_tool_ids, strict=True) + ] + loop_multimodal_fields[k] = selected + else: + loop_multimodal_fields = {} + + # Generate new completions after tool execution (using concatenated IDs, no re-tokenization) + post_tool_ids, post_tool_logprobs = self._generate_single_turn( + prompt_completion_tool_ids, loop_images, loop_multimodal_fields + ) + + # Truncate so that pct[len(prompt_ids[idx]) :] + post_tool does not exceed max_completion_length. + # The pre-regen check guarantees len(completion_tool_ids) <= max_completion_length, so any + # excess can only come from post_tool_ids. post_tool_ids is model-generated text and never + # contains image tokens, so a plain slice is safe. + for idx in range(len(idxs_with_tool)): + idx_with_tool = idxs_with_tool[idx] + completion_tool_length = len(prompt_completion_tool_ids[idx]) - len(prompt_ids[idx_with_tool]) + excess_length = completion_tool_length + len(post_tool_ids[idx]) - self.max_completion_length + if excess_length > 0: + new_len = len(post_tool_ids[idx]) - excess_length + post_tool_ids[idx] = post_tool_ids[idx][:new_len] + if logprobs is not None: + post_tool_logprobs[idx] = post_tool_logprobs[idx][:new_len] + + # Update tool_mask: the tool result should be 0 and the post-tool 1 + for idx in range(len(idxs_with_tool)): + idx_with_tool = idxs_with_tool[idx] + prompt_completion_tool_length = len(prompt_completion_tool_ids[idx]) + prompt_length = len(prompt_ids[idx_with_tool]) + completion_length = len(completion_ids[idx_with_tool]) + post_tool_length = len(post_tool_ids[idx]) + tool_length = prompt_completion_tool_length - prompt_length - completion_length + tool_mask[idx_with_tool] += [0] * tool_length + [1] * post_tool_length + if logprobs is not None: + logprobs[idx_with_tool] += [0.0] * tool_length + post_tool_logprobs[idx] + + # Update completion_ids with the new completions (after tool execution) + for idx in range(len(idxs_with_tool)): + idx_with_tool = idxs_with_tool[idx] + prompt_length = len(prompt_ids[idx_with_tool]) + pct = prompt_completion_tool_ids[idx] # = prompt-completion-tool + completion_ids[idx_with_tool] = pct[prompt_length:] + post_tool_ids[idx] + + # Decode post-tool completions. + post_tool_completions = [parse_response(self._tokenizer, ids) if ids else {} for ids in post_tool_ids] + + # Add post-tool completions to the existing completions + for idx in range(len(idxs_with_tool)): + idx_with_tool = idxs_with_tool[idx] + if post_tool_completions[idx]: # {} if post-tool completions completely truncated + completions[idx_with_tool].append(post_tool_completions[idx]) + + # Check for further tool calls + tool_calls = [completion.get("tool_calls") for completion in post_tool_completions] + idxs_with_tool = [idx for idx, tool_call in zip(idxs_with_tool, tool_calls, strict=True) if tool_call] + tool_calls = [tool_call for tool_call in tool_calls if tool_call] + iteration_num += 1 + + return tool_mask, completions, completion_ids, logprobs, tool_call_count, tool_failure_count, tool_images + + def _generate(self, prompts: list): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + # Copy the prompts to avoid modifying the original list + prompts = copy.deepcopy(prompts) + + if self.rollout_func is not None: + # Keep vLLM weights in sync for custom rollouts that rely on vLLM utilities. + if self.use_vllm and self.state.global_step != self._last_loaded_step: + if not getattr(getattr(self.vllm_generation, 'llm', None), 'shared_weights', False): + with profiling_context(self, 'sync_weights'): + self.vllm_generation.sync_weights() + self._last_loaded_step = self.state.global_step + + # Pass prompts to rollout_func preserving structured messages. + # Chat templating must happen inside rollout_func, at the backend boundary, so that + # multimodal content (images, typed content blocks) is not lost before rollout logic runs. + output = self.rollout_func(prompts, self) + required_keys = {"prompt_ids", "completion_ids", "logprobs"} + missing_keys = required_keys - output.keys() + if missing_keys: + missing_keys_list = sorted(missing_keys) + raise ValueError(f"rollout_func must return keys {missing_keys_list} in its output dict.") + extra_fields = {k: v for k, v in output.items() if k not in required_keys} + prompt_ids, completion_ids, logprobs = output["prompt_ids"], output["completion_ids"], output["logprobs"] + images = None + multimodal_fields = {} + else: + prompt_ids, images, multimodal_fields = self._tokenize_prompts(prompts) + completion_ids, logprobs = self._generate_single_turn(prompt_ids, images, multimodal_fields) + extra_fields = {} + + # Decode completions. It's important to use `parse_response` when possible, because it handles tool calls. + if is_conversational({"prompt": prompts[0]}): + if ( + Version(transformers.__version__) >= Version("5.0.0") # parse_response added in v5 + and hasattr(self._tokenizer, "response_schema") # attribute not set by default for now + and self._tokenizer.response_schema is not None # only works if the tokenizer has a schema + ): + completions = [[parse_response(self._tokenizer, ids)] for ids in completion_ids] + else: + contents = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + completions = [[{"role": "assistant", "content": content}] for content in contents] + else: + completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + + # Extract tool calls from the completions and (possibly) execute them + tool_images = [] + if self.tools: + ( + tool_mask, + completions, + completion_ids, + logprobs, + tool_call_count, + tool_failure_count, + tool_images, + ) = self._tool_call_loop( + prompts, prompt_ids, completion_ids, completions, logprobs, images, multimodal_fields + ) + # Merge tool response images into the images list for the forward pass + if any(imgs for imgs in tool_images): + if images is None: + images = [imgs if imgs else None for imgs in tool_images] + else: + images = [(existing or []) + new for existing, new in zip(images, tool_images, strict=True)] + else: + # Support custom env_mask from rollout_func (e.g., for environment feedback masking) + # Internally treated as tool_mask - marks model tokens (1) vs external tokens (0) + tool_mask = extra_fields.pop("env_mask", None) + + # Get completion length per sequence, used for logging + prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) + if tool_mask is not None: # count only model-generated tokens (tool_mask=1) + completion_lengths = torch.tensor([sum(mask) for mask in tool_mask], device=device) + else: + completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) + agg_prompt_lengths = self.accelerator.gather(prompt_lengths) + agg_completion_lengths = self.accelerator.gather(completion_lengths) + total_prompt_tokens = agg_prompt_lengths.sum() + total_completion_tokens = agg_completion_lengths.sum() # = num_items_in_batch, required for the DAPO loss + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += (total_prompt_tokens + total_completion_tokens).item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + eos_and_pad = [self._tokenizer.eos_token_id, self._tokenizer.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) + agg_is_truncated = self.accelerator.gather(is_truncated) + self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) + term_completion_lengths = agg_completion_lengths[~agg_is_truncated] + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + if self.tools: + agg_tool_call_count = self.accelerator.gather(torch.tensor(tool_call_count, device=device)).sum() + tool_call_frequency = (agg_tool_call_count / len(agg_prompt_lengths)).item() + self._metrics[mode]["tools/call_frequency"].append(tool_call_frequency) + agg_tool_failure_count = self.accelerator.gather(torch.tensor(tool_failure_count, device=device)).sum() + failure_frequency = ( + (agg_tool_failure_count / agg_tool_call_count).item() if agg_tool_call_count > 0 else 0.0 + ) + self._metrics[mode]["tools/failure_frequency"].append(failure_frequency) + + return ( + prompt_ids, + completion_ids, + tool_mask, + completions, + total_completion_tokens, + logprobs, + extra_fields, + images, + tool_images, + ) + + def _generate_and_score_completions( + self, inputs: list[dict[str, torch.Tensor | Any]] + ) -> dict[str, torch.Tensor | Any]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + # Unsloth: Extract per-sample chat_template_kwargs before metadata is lost + _ct_ = getattr(self.processing_class, 'chat_template', None) or '' + _sk_ = {'prompt', 'chosen', 'rejected', 'completion', 'messages', 'label', + 'images', 'image', 'videos', 'video', 'audios', 'audio'} + self._unsloth_batch_chat_kwargs = [] + for _inp_ in inputs: + _kw_ = {} + if isinstance(_inp_, dict): + for _k_ in _inp_.keys() - _sk_: + if _k_ in _ct_ and isinstance(_inp_[_k_], str): + _kw_[_k_] = _inp_[_k_] + self._unsloth_batch_chat_kwargs.append(_kw_) + if self.environments: + for prompt, environment, reset_kwargs in zip(prompts, self.environments, inputs, strict=True): + observation = environment.reset(**reset_kwargs) + if observation is None: + continue + if isinstance(observation, list) and isinstance(prompt[-1]["content"], str): + prompt[-1]["content"] = [{"type": "text", "text": prompt[-1]["content"]}] + if isinstance(observation, str) and isinstance(prompt[-1]["content"], list): + observation = [{"type": "text", "text": observation}] + prompt[-1]["content"] += observation + + if "images" in inputs[0]: + images = [example.get("images") for example in inputs] + elif "image" in inputs[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] + else: + images = None + # Transformers requires at least one image in the batch, otherwise it throws an error + if images is not None and all(img_list == [] for img_list in images): + images = None + + # If the prompts are conversational and the inputs contain images, we need to convert the prompts from + # [{"role": "user", "content": "What color is the sky?"}] to + # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] + if images is not None: + if not is_conversational(inputs[0]): + raise ValueError( + "Multimodal training requires conversational prompts. It looks like the dataset contains " + "non-conversational inputs, likely because a chat template was applied before passing the dataset " + "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat " + "template internally." + ) + prompts = [ + prepare_multimodal_messages(prompt, images=image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ] + + dataset_images = images # preserve dataset images before _generate may overwrite + ( + prompt_ids_list, + completion_ids_list, + tool_mask_list, + completions, + num_items_in_batch, + sampling_per_token_logps_list, + extra_fields, + images, + tool_images, + ) = self._generate(prompts) + + _unsloth_clear_stateful_mrope( + self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False) + ) + if images is None: + images = dataset_images # restore dataset images (rollout_func path returns None) + + # Convert lists of token IDs to padded tensors + prompt_ids = [torch.tensor(ids) for ids in prompt_ids_list] + prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] + prompt_ids = pad( + prompt_ids, + padding_value=self._tokenizer.pad_token_id, + padding_side="left", + pad_to_multiple_of=self.pad_to_multiple_of, + ).to(device=device) + prompt_mask = pad( + prompt_mask, padding_value=0, padding_side="left", pad_to_multiple_of=self.pad_to_multiple_of + ).to(device=device) + completion_ids = [torch.tensor(ids) for ids in completion_ids_list] + completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] + completion_ids = pad( + completion_ids, + padding_value=self._tokenizer.pad_token_id, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ).to(device=device) + completion_mask = pad( + completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ).to(device=device) + if sampling_per_token_logps_list is not None: + sampling_per_token_logps = [torch.tensor(logps) for logps in sampling_per_token_logps_list] + sampling_per_token_logps = pad( + sampling_per_token_logps, + padding_value=0.0, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ).to(device=device) + else: + sampling_per_token_logps = None + if tool_mask_list is not None: + tool_mask = [torch.tensor(mask) for mask in tool_mask_list] + tool_mask = pad( + tool_mask, padding_value=1, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ).to(device=device) + else: + tool_mask = None + + # If mask_truncated_completions is enabled, zero out truncated completions for attention and loss masking + if self.mask_truncated_completions: + eos_and_pad = [self._tokenizer.eos_token_id, self._tokenizer.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) + # Mask completion_mask for attention masking + completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() + # Also mask tool_mask for consistency in multi-turn training + if tool_mask is not None: + tool_mask = tool_mask * (~is_truncated).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + max_left_pad = None + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + try: + # TRL 0.23.1 and below path + if not has_images: + # Left pad prompt before calculation old and ref hidden states + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() + except: + # TRL 0.24.0 and below path + if images is None: + # Left pad prompt before calculation old and ref hidden states + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() + self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True)) + + num_images = [len(img_list) if img_list else 0 for img_list in images] if images is not None else None + + # Get forward_kwargs for models with multimodal inputs. + # When tool images are present (from _tool_call_loop), use image_processor directly and build + # mm_token_type_ids from prompt_completion_ids. Otherwise, use the full processor pipeline + # which returns model-specific keys (image_sizes, pixel_attention_mask, etc.). + if self.tools and any(imgs for imgs in tool_images) and self._is_vlm: + flat_images = [img for img_list in images if img_list for img in img_list] + image_inputs = self.processing_class.image_processor(images=flat_images, return_tensors="pt") + image_inputs = super()._prepare_inputs(image_inputs) + forward_kwargs = dict(image_inputs) + elif images is not None: + prompts_text = [ + apply_chat_template( + {"prompt": prompt}, self.processing_class, tools=self.tools, **self.chat_template_kwargs + )["prompt"] + for prompt in prompts + ] + prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") + prompt_inputs = super()._prepare_inputs(prompt_inputs) + forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} + else: + forward_kwargs = {} + + # Recover LFM2-VL tile counts; the full processor drops row/column metadata. + num_tiles = None + if images is not None and "spatial_shapes" in forward_kwargs: + image_info = self.processing_class.image_processor( + images=images, return_tensors="pt", return_row_col_info=True + ) + tiles_per_image = image_info["image_rows"] * image_info["image_cols"] + if self.processing_class.image_processor.use_thumbnail: + tiles_per_image = tiles_per_image + (tiles_per_image > 1).to(tiles_per_image.dtype) + num_tiles = [group.sum().item() for group in torch.split(tiles_per_image, num_images)] + + # If token_type_ids are used, extend them with zeros for the completion part + if "token_type_ids" in forward_kwargs: + token_type_ids = forward_kwargs["token_type_ids"] + if self.pad_to_multiple_of is not None: + # Needed only with pad_to_multiple_of: otherwise prompt_ids and token_type_ids must have equal len + padding_size = prompt_ids.size(1) - token_type_ids.size(1) + if padding_size > 0: + token_type_ids = torch.cat( + [token_type_ids.new_zeros((token_type_ids.size(0), padding_size)), token_type_ids], dim=1 + ) + forward_kwargs["token_type_ids"] = torch.cat( + [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + # If mm_token_type_ids are used, extend them with zeros for the completion part + if "mm_token_type_ids" in forward_kwargs: + mm_token_type_ids = forward_kwargs["mm_token_type_ids"] + if self.pad_to_multiple_of is not None: + # Needed only with pad_to_multiple_of: otherwise prompt_ids and mm_token_type_ids must have equal len + padding_size = prompt_ids.size(1) - mm_token_type_ids.size(1) + if padding_size > 0: + mm_token_type_ids = torch.cat( + [mm_token_type_ids.new_zeros((mm_token_type_ids.size(0), padding_size)), mm_token_type_ids], + dim=1, + ) + forward_kwargs["mm_token_type_ids"] = torch.cat( + [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + if "mm_token_type_ids" in forward_kwargs or "image_grid_thw" in forward_kwargs: + _mm_token_type_ids = _unsloth_fix_mm_token_type_ids( + self.processing_class, + prompt_completion_ids, + forward_kwargs.get("mm_token_type_ids", None), + completion_ids = completion_ids, + ) + if _mm_token_type_ids is not None: + forward_kwargs["mm_token_type_ids"] = _mm_token_type_ids + + # For VLM tool images: build token type IDs from the full prompt_completion_ids. + # This must happen AFTER the token_type_ids/mm_token_type_ids extension blocks above, + # because our version already covers the full sequence (images are in the completion, + # not just the prompt). + if self.tools and any(imgs for imgs in tool_images) and self._is_vlm: + mm_ids = torch.zeros_like(prompt_completion_ids) + if self._image_pad_token_id is not None: + mm_ids[prompt_completion_ids == self._image_pad_token_id] = 1 + if self._video_pad_token_id is not None: + mm_ids[prompt_completion_ids == self._video_pad_token_id] = 2 + + # Use the same key the model expects: token_type_ids for models like Gemma, + # mm_token_type_ids for models like Qwen. + image_grid_thw = forward_kwargs.get("image_grid_thw") + if image_grid_thw is not None: + forward_kwargs["mm_token_type_ids"] = mm_ids + else: + forward_kwargs["token_type_ids"] = mm_ids + + # Truncation safety (Qwen-style models with image_grid_thw only): if + # max_completion_length truncated some image tokens, the number of image pad tokens + # in input_ids won't match pixel_values features. Check per-sample and drop ALL + # images for any sample with a mismatch (safe fallback). + if image_grid_thw is not None and num_images is not None: + merge_length = getattr(self.processing_class.image_processor, "merge_size", 2) ** 2 + img_offset = 0 + has_mismatch = False + for b in range(mm_ids.shape[0]): + sample_tokens = (mm_ids[b] == 1).sum().item() + sample_features = 0 + for i in range(num_images[b]): + grid_idx = img_offset + i + if grid_idx < image_grid_thw.shape[0]: + sample_features += image_grid_thw[grid_idx].prod().item() // merge_length + if sample_tokens != sample_features: + has_mismatch = True + break + img_offset += num_images[b] + + if has_mismatch: + # Drop all images: safer than partial trim which is error-prone + forward_kwargs.pop("pixel_values", None) + forward_kwargs.pop("image_grid_thw", None) + mm_ids.zero_() + forward_kwargs["mm_token_type_ids"] = mm_ids + num_images = None + + # When gradient checkpointing is enabled with use_reentrant=True (non default), calling the model inside a + # torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True"). + # Temporarily disable checkpointing to avoid this warning during inference. + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + # If the generation and optimization steps are misaligned—i.e., if generation does not occur at the end of + # a full optimizer step (when gradient_accumulation_steps is not a multiple of generate_every)—then the + # samples may come from an earlier version of the model. In that case, we need to track old_per_token_logps + # for importance sampling. If the steps are aligned, importance sampling isn't necessary and we set + # old_per_token_logps to None. + # When using vLLM, we always compute old_per_token_logps for importance sampling, it was shown that the + # distribution mismatch between vLLM and the training model can be large and harm the training. + generate_every = self.args.steps_per_generation * self.num_iterations # generation frequency + + if self.args.gradient_accumulation_steps % generate_every != 0 or ( + self.use_vllm + ): + old_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size, + num_images=num_images, + num_tiles=num_tiles, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids + ) + else: + old_per_token_logps = None + + # Compute the importance sampling ratio when using vLLM, to correct for potential distribution mismatch + if False and self.use_vllm and self.vllm_importance_sampling_correction: + mask = completion_mask if tool_mask is None else completion_mask * tool_mask + per_token_logps_diff = (old_per_token_logps - sampling_per_token_logps) * mask + + sequence_level_is = self.vllm_importance_sampling_mode in ["sequence_mask", "sequence_truncate"] + if sequence_level_is: + per_sequence_logps_diff = per_token_logps_diff.sum(dim=-1, keepdim=True) + logps_diff = per_sequence_logps_diff + else: + logps_diff = per_token_logps_diff + + vllm_importance_sampling_ratio = torch.exp(logps_diff) + + # vllm_importance_sampling_ratio.shape: + # token_* modes: (B, T) (per-token ratio) + # sequence_* modes: (B, 1) (per-sequence ratio) + + if self.vllm_importance_sampling_mode in ["sequence_truncate", "token_truncate"]: + vllm_importance_sampling_ratio = torch.clamp( + vllm_importance_sampling_ratio, + min=self.vllm_importance_sampling_clip_min, + max=self.vllm_importance_sampling_clip_max, + ) + elif self.vllm_importance_sampling_mode in ["sequence_mask", "token_mask"]: + min_val = ( + self.vllm_importance_sampling_clip_min + if self.vllm_importance_sampling_clip_min is not None + else -math.inf + ) + max_val = ( + self.vllm_importance_sampling_clip_max + if self.vllm_importance_sampling_clip_max is not None + else math.inf + ) + + invalid_mis_mask = (vllm_importance_sampling_ratio < min_val) | ( + vllm_importance_sampling_ratio > max_val + ) + vllm_importance_sampling_ratio = vllm_importance_sampling_ratio.masked_fill( + invalid_mis_mask, value=0.0 + ) + else: + raise ValueError( + f"Unknown vLLM importance sampling level: {self.vllm_importance_sampling_mode}. Possible values are 'token_truncate', 'token_mask', 'sequence_truncate', and 'sequence_mask'." + ) + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( + self.ref_model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + num_tiles=num_tiles, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids + ) + else: + # When training a PEFT adapter, how we obtain the reference depends on the setup: + # - New adapter: disabling adapters yields the base model. + # - Re-training an existing adapter: an initial copy is loaded under the name "ref". + model = self.accelerator.unwrap_model(self.model) + with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): + ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + num_tiles=num_tiles, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids + ) + else: + ref_per_token_logps = None + + # Decode + prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + + # Merge extra_fields from rollout_func into inputs for reward functions + if extra_fields: + for i, inp in enumerate(inputs): + for key, values in extra_fields.items(): + if isinstance(values, list) and i < len(values): + inp[key] = values[i] + elif not isinstance(values, list): + inp[key] = values + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + if images is not None: + rewards_per_func = self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list) + else: + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + num_generations = self.num_generations if mode == "train" else self.num_generations_eval + + # A completion for which every reward function returned None is unscorable. nansum would collapse it to 0, + # which both biases the per-group baseline and hands the completion a spurious advantage. Mark these rows NaN + # so they're excluded from the (nan-aware) baseline below; their advantage is forced to 0 afterwards. + unscorable_mask = torch.isnan(rewards_per_func).all(dim=1) + + if self.multi_objective_aggregation == "sum_then_normalize": + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + rewards[unscorable_mask] = torch.nan + mean_grouped_rewards = torch.nanmean(rewards.view(-1, num_generations), dim=1) + mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(num_generations, dim=0) + if self.scale_rewards in ["group", "none"]: + # If self.scale_rewards = "none", we'll only use std_rewards to check for zero std for logging + if num_generations > 1: + std_rewards = nanstd(rewards.view(-1, num_generations), dim=1) + std_rewards = std_rewards.repeat_interleave(num_generations, dim=0) + else: # doesn't occur during training, but could occur in eval when num_generations_eval=1 + std_rewards = torch.zeros_like(rewards) + elif self.scale_rewards == "batch": + # Compute global std + if rewards.numel() > 1: + std_rewards = nanstd(rewards).expand_as(rewards) + else: # doesn't occur during training, but could occur in eval when num_generations_eval=batch_size=1 + std_rewards = torch.zeros_like(rewards) + else: + raise ValueError( + f"Invalid value for scale_rewards: {self.scale_rewards}. Must be one of 'batch', 'group', or 'none'." + ) + + advantages = rewards - mean_grouped_rewards + if self.scale_rewards != "none": + advantages = advantages / (std_rewards + 1e-4) + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) # for logging + + elif self.multi_objective_aggregation == "normalize_then_sum": + grouped = rewards_per_func.view(-1, num_generations, len(self.reward_funcs)) + mean_k = torch.nanmean(grouped, dim=1, keepdim=True) + std_k = nanstd(grouped, dim=1, keepdim=True) if num_generations > 1 else torch.zeros_like(mean_k) + reward_k = (grouped - mean_k) / (std_k + 1e-4) + reward_k = reward_k.view(-1, len(self.reward_funcs)) + rewards = (reward_k * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + rewards[unscorable_mask] = torch.nan + std_rewards = nanstd(rewards).expand_as(rewards) if rewards.numel() > 1 else torch.zeros_like(rewards) + advantages = (rewards - torch.nanmean(rewards)) / (std_rewards + 1e-4) + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) # for logging + + else: + raise ValueError( + f"Invalid multi_objective_aggregation: {self.multi_objective_aggregation}. Must be " + "'sum_then_normalize' or 'normalize_then_sum'." + ) + + # Unscorable completions (every reward func returned None) carry no learning signal: their reward is NaN here, + # so zero their advantage to keep them from moving the policy. + advantages = torch.nan_to_num(advantages, nan=0.0) + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_func_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) + rewards = (rewards_per_func * self.reward_weights.to(rewards_per_func.device).unsqueeze(0)).nansum(dim=1) + rewards[unscorable_mask] = torch.nan # exclude unscorable rows from the logged reward stats + self._metrics[mode]["reward"].append(torch.nanmean(rewards).item()) + self._metrics[mode]["reward_std"].append(nanstd(rewards).item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._logs["prompt"].extend(gather_object(prompts_text)) + self._logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._logs["advantages"].extend(all_process_advantages.tolist()) + + # Flush user-logged extra columns (from log_extra), gathering across processes. + # Keys must be sorted so that all ranks call gather_object in the same order, otherwise values + # get mis-attributed across columns (dict insertion order may differ between processes). + for column in sorted(self._pending_extra_logs): + self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column])) + self._pending_extra_logs.clear() + + # Flush user-logged metrics (from log_metric), averaging across processes. + # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values + # get mis-attributed across metrics (dict insertion order may differ between processes). + for name in sorted(self._pending_metrics): + values = self._pending_metrics[name] + local_mean = sum(values) / len(values) + global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item() + self._metrics[mode][name].append(global_mean) + self._pending_metrics.clear() + + if images is not None: + self._logs["images"].extend(gather_object(images)) + + if False and self.use_vllm and self.vllm_importance_sampling_correction: + delta = torch.abs(old_per_token_logps - sampling_per_token_logps) + mask = completion_mask.bool() if tool_mask is None else (completion_mask * tool_mask).bool() + delta = delta[mask] + mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( + self.accelerator.gather(mean_delta).mean().item() + ) + self._metrics[mode]["sampling/sampling_logp_difference/max"].append( + self.accelerator.gather(max_delta).max().item() + ) + if sequence_level_is: + flat_is_ratio = vllm_importance_sampling_ratio.flatten() + else: + flat_is_ratio = vllm_importance_sampling_ratio[mask] + + min_importance_sampling_ratio = ( + torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + mean_importance_sampling_ratio = ( + torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + max_importance_sampling_ratio = ( + torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( + nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( + self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( + nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + ) + + output = { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "advantages": advantages, + "num_items_in_batch": num_items_in_batch, + } + if old_per_token_logps is not None: + output["old_per_token_logps"] = old_per_token_logps + if False and self.use_vllm and self.vllm_importance_sampling_correction: + output["importance_sampling_ratio"] = vllm_importance_sampling_ratio + if sampling_per_token_logps is not None: + output["sampling_per_token_logps"] = sampling_per_token_logps + if ref_per_token_logps is not None: + output["ref_per_token_logps"] = ref_per_token_logps + if "pixel_values" in forward_kwargs: + output["pixel_values"] = forward_kwargs["pixel_values"] + if "image_grid_thw" in forward_kwargs: + output["image_grid_thw"] = forward_kwargs["image_grid_thw"] + if "pixel_attention_mask" in forward_kwargs: + output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] + if "spatial_shapes" in forward_kwargs: + output["spatial_shapes"] = forward_kwargs["spatial_shapes"] + if "image_sizes" in forward_kwargs: + output["image_sizes"] = forward_kwargs["image_sizes"] + if "token_type_ids" in forward_kwargs: + output["token_type_ids"] = forward_kwargs["token_type_ids"] + if "mm_token_type_ids" in forward_kwargs: + output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"] + if "image_position_ids" in forward_kwargs: + output["image_position_ids"] = forward_kwargs["image_position_ids"] + if images is not None: + output["num_images"] = num_images + if max_left_pad is not None: + output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1) + try: + if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False): + output["sampling_per_token_logps"] = sampling_per_token_logps + except NameError: + output["sampling_per_token_logps"] = None + if num_tiles is not None: + output["num_tiles"] = num_tiles + if tool_mask is not None: + output["tool_mask"] = tool_mask + return output + + def compute_liger_loss(self, unwrapped_model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # Get the last hidden state of the model + last_hidden_state = self._get_last_hidden_state( + unwrapped_model, + input_ids, + attention_mask, + logits_to_keep, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + inputs.get("pixel_attention_mask"), + inputs.get("spatial_shapes"), + inputs.get("image_sizes"), + inputs.get("image_position_ids"), + ) + + # Apply tool_mask (from env_mask) for loss computation in multi-turn training scenarios + loss_mask = completion_mask if "tool_mask" not in inputs else completion_mask * inputs["tool_mask"] + lm_head_weight = unwrapped_model.lm_head.weight + lm_head_bias = unwrapped_model.lm_head.bias + # Liger reads `lm_head` directly instead of through `model.forward()`, so its ZeRO-3 gather hook never fires + # and the fused matmul gets an empty shard. Gather the weight/bias ourselves for the call (the weight grad is + # computed during this forward, so it isn't needed in the backward). Skip it when already gathered: with tied + # embeddings `embed_tokens` keeps the weight `AVAILABLE`, and re-partitioning on exit breaks its tracking. + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + gather_ctx = nullcontext() + if deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3: + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + + params = [lm_head_weight] if lm_head_bias is None else [lm_head_weight, lm_head_bias] + if any(p.ds_status != ZeroParamStatus.AVAILABLE for p in params): + import deepspeed + + gather_ctx = deepspeed.zero.GatheredParameters(params, modifier_rank=None) + with gather_ctx: + loss, metrics = self.liger_grpo_loss( + _input=last_hidden_state, + lin_weight=lm_head_weight, + selected_token_ids=completion_ids, + # The attention_mask parameter in liger loss is actually used as a loss mask (not model attention) + attention_mask=loss_mask, + advantages=inputs["advantages"], + bias=lm_head_bias, + old_per_token_logps=inputs.get("old_per_token_logps"), + ref_per_token_logps=inputs.get("ref_per_token_logps"), + vllm_is_ratio=inputs.get("importance_sampling_ratio"), + ) + # Extract metrics from the liger_grpo_loss output + # KL divergence is the first metric when beta is non-zero + mean_kl = metrics[0] if self.beta != 0.0 else None + clip_ratio = metrics[-1] + + mode = "train" if self.model.training else "eval" + if self.beta != 0.0: + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).mean().item()) + self._metrics[mode]["clip_ratio"].append(self.accelerator.gather(clip_ratio).mean().item()) + normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval + return loss / normalizer + + def compute_loss( + self, + model, + inputs, + return_outputs = False, + num_items_in_batch = None, + ): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + # Compute the per-token log probabilities for the model + + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = ( + inputs["completion_ids"], + inputs["completion_mask"], + ) + pixel_values, image_grid_thw = ( + inputs.get("pixel_values", None), + inputs.get("image_grid_thw", None), + ) + pixel_attention_mask, image_sizes = ( + inputs.get("pixel_attention_mask", None), + inputs.get("image_sizes", None), + ) + num_images = inputs.get("num_images", None) + # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models + token_type_ids = inputs.get("token_type_ids", None) + mm_token_type_ids = inputs.get("mm_token_type_ids", None) + num_items_in_batch = inputs.get("num_items_in_batch", None) + sampling_per_token_logps = inputs.get("sampling_per_token_logps", None) + tool_mask = inputs.get("tool_mask", None) + # Missing when evaluate() runs standalone; eval does not accumulate, so + # fall back to 1 to avoid underreporting eval_loss (#2464). + current_gradient_accumulation_steps = getattr( + self, "current_gradient_accumulation_steps", 1 + ) + num_processes = self.accelerator.num_processes + + input_ids = torch.cat([prompt_ids, completion_ids], dim = 1) + bsz, qlen = input_ids.shape + attention_mask = torch.cat([prompt_mask, completion_mask], dim = 1) + if mm_token_type_ids is not None or image_grid_thw is not None: + mm_token_type_ids = _unsloth_fix_mm_token_type_ids( + self.processing_class, + input_ids, + mm_token_type_ids, + completion_ids = completion_ids, + ) + # attention_mask = None + logits_to_keep = completion_ids.size( + 1 + ) # we only need to compute the logits for the completion tokens + _input_ids = input_ids + _logits_to_keep = logits_to_keep + + get_logps_func = ( + lambda model, + input_ids, + attention_mask, + logits_to_keep, + batch_size = None, + compute_entropy = False, + compute_efficient = False: self._get_per_token_logps( + model, input_ids, attention_mask, logits_to_keep, compute_efficient + ) + 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, + compute_efficient, + )[0] + ) # logps + + per_token_logps = get_logps_func( + model, input_ids, attention_mask, logits_to_keep, compute_efficient = True + ) + # Compute the KL divergence between the model and the reference model + # _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves. + # https://github.com/huggingface/trl/blob/05bc43e960396581e458195b8388efe6b82cae1f/trl/trainer/grpo_trainer.py#L1328 + # if self.beta != 0.0: + # with torch.inference_mode(), model.disable_adapter(): + # ref_per_token_logps = per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep) + # else: + # ref_per_token_logps = None + ref_logps = inputs.get("ref_per_token_logps", None) + # per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + # x - x.detach() allows for preserving gradients from x + advantages = inputs["advantages"] + # per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + # per_token_loss = -(per_token_loss - self.beta * per_token_kl) + # loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + old_logps = inputs.get("old_per_token_logps", None) + + input_ids = input_ids[:, -logits_to_keep:] + + # Get logit softcapping and logit scale + logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma + logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere + if logit_scale_multiply is None: + logit_scale_multiply = 0 + logit_scale_divide = getattr(model.config, "logits_scaling", 0) # Granite + if logit_scale_divide is None: + logit_scale_divide = 0 + + max_left_pad = inputs.get("max_left_pad", 0) + if per_token_logps is not None: + loss_mask = completion_mask + if tool_mask is not None: + if tool_mask.shape != completion_mask.shape: + raise ValueError( + "tool_mask/env_mask must have the same shape as completion_mask" + ) + loss_mask = completion_mask * tool_mask.to( + device = completion_mask.device, + dtype = completion_mask.dtype, + ) + ( + loss, + completion_length, + mean_kl, + delta, + flat_is_ratio, + coef_1, + completion_mask, + ) = grpo_compute_loss_slow( + ref_logps, + per_token_logps, + old_logps, + sampling_per_token_logps, + input_ids, + loss_mask, + self.beta, + advantages, + pixel_values = pixel_values, + image_grid_thw = image_grid_thw, + loss_type = self.args.loss_type, + importance_sampling_level = self.importance_sampling_level, + epsilon_low = self.epsilon_low, + epsilon_high = self.epsilon_high, + max_completion_length = self.args.max_completion_length, + delta = self.args.delta, + temperature = self.args.temperature, + max_left_pad = max_left_pad, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + num_items_in_batch = num_items_in_batch, + current_gradient_accumulation_steps = current_gradient_accumulation_steps, + num_processes = num_processes, + ) + else: + + def _unsloth_requires_multi_image_zoo(value): + if value is None: + return False + if isinstance(value, torch.Tensor): + counts = value.detach().cpu().reshape(-1).tolist() + else: + counts = list(value) + return any(int(n) != 1 for n in counts) + + if _unsloth_requires_multi_image_zoo(num_images) and not getattr( + self, "_unsloth_grpo_zoo_checked", False + ): + _supports_num_images = ( + "num_images" in inspect.signature(grpo_accumulated_loss).parameters + ) + if not _supports_num_images: + try: + _zoo_src = inspect.getsource(grpo_accumulated_loss) + except (TypeError, OSError): + _zoo_src = "" + _supports_num_images = "num_images" in _zoo_src + if not _supports_num_images: + raise RuntimeError( + "Multi-image GRPO requires an unsloth_zoo build whose " + "grpo_accumulated_loss handles num_images. Please upgrade " + "unsloth_zoo (see https://github.com/unslothai/unsloth-zoo/pull/613)." + ) + self._unsloth_grpo_zoo_checked = True + if tool_mask is not None and not getattr( + self, "_unsloth_grpo_tool_mask_zoo_checked", False + ): + _supports_tool_mask = ( + "tool_mask" in inspect.signature(grpo_accumulated_loss).parameters + ) + if not _supports_tool_mask: + try: + _zoo_src = inspect.getsource(grpo_accumulated_loss) + except (TypeError, OSError): + _zoo_src = "" + _supports_tool_mask = "tool_mask" in _zoo_src + if not _supports_tool_mask: + raise RuntimeError( + "env_mask/tool_mask GRPO requires an unsloth_zoo build whose " + "grpo_accumulated_loss handles tool_mask. Please upgrade " + "unsloth_zoo." + ) + self._unsloth_grpo_tool_mask_zoo_checked = True + _grpo_accumulated_loss_kwargs = {} + if tool_mask is not None: + _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask + if hasattr(self.args, "loss_type"): + ( + loss, + completion_length, + mean_kl, + delta, + flat_is_ratio, + coef_1, + completion_mask, + ) = grpo_accumulated_loss( + trainer = self, + input_ids = _input_ids, + pixel_values = pixel_values, + image_grid_thw = image_grid_thw, + pixel_attention_mask = pixel_attention_mask, + image_sizes = image_sizes, + num_images = num_images, + logits_to_keep = logits_to_keep, + completion_mask = completion_mask, + advantages = advantages, + old_logps = old_logps, + ref_logps = ref_logps, + n_chunks = self.args.unsloth_num_chunks, + loss_type = self.args.loss_type, + importance_sampling_level = self.importance_sampling_level, + epsilon_low = self.epsilon_low, + epsilon_high = self.epsilon_high, + max_completion_length = self.args.max_completion_length, + delta = self.args.delta, + temperature = self.args.temperature, + max_left_pad = max_left_pad, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + attention_mask = attention_mask, + num_items_in_batch = num_items_in_batch, + current_gradient_accumulation_steps = current_gradient_accumulation_steps, + num_processes = num_processes, + sampling_per_token_logps = sampling_per_token_logps, + token_type_ids = token_type_ids, + mm_token_type_ids = mm_token_type_ids, + **_grpo_accumulated_loss_kwargs, + ) + else: + # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 + loss, completion_length, mean_kl, coef_1, completion_mask = grpo_accumulated_loss( + trainer = self, + input_ids = _input_ids, + pixel_values = pixel_values, + image_grid_thw = image_grid_thw, + pixel_attention_mask = pixel_attention_mask, + image_sizes = image_sizes, + num_images = num_images, + logits_to_keep = logits_to_keep, + completion_mask = completion_mask, + advantages = advantages, + old_logps = old_logps, + ref_logps = ref_logps, + n_chunks = self.args.unsloth_num_chunks, + temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + attention_mask = attention_mask, + token_type_ids = token_type_ids, + mm_token_type_ids = mm_token_type_ids, + **_grpo_accumulated_loss_kwargs, + ) + if "train" in self._metrics: + mode = "eval" if self.control.should_evaluate else "train" + self._metrics[mode]["completion_length"].append(completion_length.item()) + self._metrics[mode]["kl"].append(mean_kl.item()) + else: + self._metrics["completion_length"].append(completion_length.item()) + self._metrics["kl"].append(mean_kl.item()) + + if ( + self.use_vllm + and delta is not None + and getattr(self, "vllm_importance_sampling_correction", False) + ): + mean_delta = ( + torch.mean(delta) + if delta.numel() > 0 + else torch.tensor(0.0, device = self.model.device) + ) + max_delta = ( + torch.max(delta) + if delta.numel() > 0 + else torch.tensor(0.0, device = self.model.device) + ) + self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( + self.accelerator.gather(mean_delta).mean().item() + ) + self._metrics[mode]["sampling/sampling_logp_difference/max"].append( + self.accelerator.gather(max_delta).max().item() + ) + + min_importance_sampling_ratio = ( + torch.min(flat_is_ratio) + if flat_is_ratio.numel() > 0 + else torch.tensor(0.0, device = self.model.device) + ) + mean_importance_sampling_ratio = ( + torch.mean(flat_is_ratio) + if flat_is_ratio.numel() > 0 + else torch.tensor(0.0, device = self.model.device) + ) + max_importance_sampling_ratio = ( + torch.max(flat_is_ratio) + if flat_is_ratio.numel() > 0 + else torch.tensor(0.0, device = self.model.device) + ) + self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( + self.accelerator.gather(min_importance_sampling_ratio) + .nan_to_num(nan = float("inf")) + .min() + .item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( + self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( + self.accelerator.gather(max_importance_sampling_ratio) + .nan_to_num(nan = float("-inf")) + .max() + .item() + ) + + completion_token_count = completion_mask.sum().clamp(min = 1.0) + + def masked_batch_mean(x): + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return x.mean() + else: + return (x * completion_mask).sum() / completion_token_count + + if advantages.dim() == 1: + advantages = advantages.unsqueeze(1) + + if self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo"]: + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) + is_region_clipped = is_low_clipped | is_high_clipped + + low_clip = masked_batch_mean(is_low_clipped.float()) + high_clip = masked_batch_mean(is_high_clipped.float()) + clip_ratio = masked_batch_mean(is_region_clipped.float()) + + gathered_low_clip = self.accelerator.gather(low_clip) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(high_clip) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(clip_ratio) + self._metrics[mode]["clip_ratio/region_mean"].append( + gathered_clip_ratio.nanmean().item() + ) + elif self.loss_type == "cispo": + is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0) + cispo_clip_ratio = masked_batch_mean(is_cispo_clipped.float()) + gathered_cispo_clip_ratio = self.accelerator.gather(cispo_clip_ratio) + self._metrics[mode]["cispo_clip_ratio"].append( + gathered_cispo_clip_ratio.nanmean().item() + ) + + return loss + + @staticmethod + def get_off_policy_mask( + advantages: torch.Tensor, + per_token_logps: torch.Tensor, + sampling_per_token_logps: torch.Tensor, + mask: torch.Tensor, + off_policy_threshold: float, + ) -> torch.Tensor: + """ + Computes the Off-Policy Sequence Mask from DeepSeek-V3.2 paper. Returns a (B, 1) tensor where 1.0 indicates + "Keep" and 0.0 indicates "Drop". + """ + # forward KL div: log(pi_old) - log(pi_theta) + kl_div = sampling_per_token_logps - per_token_logps.detach() + # Sequence-level Mean KL (ignoring prompt+padding) + seq_kl_sum = (kl_div * mask).sum(dim=1, keepdim=True) + avg_seq_kl = seq_kl_sum / mask.sum(dim=1, keepdim=True).clamp(min=1.0) + # Keep if (Advantage >= 0) OR (KL <= delta) + is_pos_adv = advantages >= 0 + is_low_kl = avg_seq_kl <= off_policy_threshold + return (is_pos_adv | is_low_kl).to(dtype=mask.dtype) # (B, 1) + + @staticmethod + @torch.no_grad() + def get_gamma_weights( + advantages: torch.Tensor, + log_ratio_per_token: torch.Tensor, + mask: torch.Tensor, + importance_sampling_ratio: torch.Tensor | None, # (B, T) + k_pos: float = 2.0, + lambda_pos: float = 3.0, + k_neg: float = 3.0, + lambda_neg: float = 2.0, + ) -> torch.Tensor: + """ + Computes the Gamma weights for the VESPO loss. For reference: + φ(w) = e^λ × w^k × e^{-λw} is the gamma weighting (normalized so φ(1)=1) + with w = sequence-level importance sampling ratio + note: we will compute φ(w) in log space + + φ(w) is detached via @torch.no_grad(), only acts as gradient scaling coefficient + + VESPO loss = -φ(w) × A × log_prob, gradient naturally gives φ(w) × A × ∇log π + """ + # reducing clamp range directly to log(1e-8) ~ -18.42, to avoid recomputing log_w=log(w.clamp(min=1e-8)) later + # This is solely for matching truthfully the original implementation, otherwise keeping -20 could be fine. + lower_clamp = math.log(1e-8) + + # Sequence-level log ratio Σ log(π_θ/π_old) (not a mean like for `log_importance_weights`) + log_ratio_clamped = torch.clamp(log_ratio_per_token, -20.0, 20.0) + seq_log_ratio = torch.sum(log_ratio_clamped * mask, dim=-1, keepdim=True) # (B, 1) + + # Apply token-level TIS or MIS correction (in log space) + if importance_sampling_ratio is not None: + log_is_ratio = torch.clamp(torch.log(importance_sampling_ratio), lower_clamp, 20.0) + # log(w) = log(π_θ/π_old) + log(π_old/π_sampler) + seq_log_ratio += torch.sum(log_is_ratio, dim=-1, keepdim=True) + + log_w_seq = torch.clamp(seq_log_ratio, lower_clamp, 20.0) + w_seq = torch.exp(log_w_seq) + + # compute k and lambda based on advantage sign + is_nonneg_adv = advantages >= 0 + k_seq = torch.where(is_nonneg_adv, k_pos, k_neg) + lambda_seq = torch.where(is_nonneg_adv, lambda_pos, lambda_neg).clamp(min=1e-4) + + # log(φ(w)) = λ + k × log(w) - λ × w + log_phi = lambda_seq + k_seq * log_w_seq - lambda_seq * w_seq + phi_seq = torch.exp(log_phi).nan_to_num(nan=0.0, posinf=0.0, neginf=0.0) + + return phi_seq # (B, 1) + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + mask = completion_mask if "tool_mask" not in inputs else completion_mask * inputs["tool_mask"] + + # Compute the per_token_logps and the entropy at each position in the completion + per_token_logps, entropies, aux_loss = self._get_per_token_logps_and_entropies( + model, + input_ids, + attention_mask, + logits_to_keep, + compute_entropy=True, + compute_aux_loss=self.aux_loss_enabled, + pixel_values=inputs.get("pixel_values"), + image_grid_thw=inputs.get("image_grid_thw"), + num_images=inputs.get("num_images"), + pixel_attention_mask=inputs.get("pixel_attention_mask"), + spatial_shapes=inputs.get("spatial_shapes"), + num_tiles=inputs.get("num_tiles"), + image_sizes=inputs.get("image_sizes"), + token_type_ids=inputs.get("token_type_ids"), + mm_token_type_ids=inputs.get("mm_token_type_ids"), + image_position_ids=inputs.get("image_position_ids"), + ) + + if self.top_entropy_quantile < 1.0: + entropy_mask = self.get_high_entropy_mask(entropies, mask, 1 - self.top_entropy_quantile) + else: + entropy_mask = None + + # Compute the loss + advantages = inputs["advantages"] + # In the base GRPO implementation, advantages are expected to have shape (B,). To support subclasses that + # provide advantages with shape (B, T) (e.g., MiniLLM), we *conditionally* unsqueeze the tensor. + if advantages.dim() == 1: + advantages = advantages.unsqueeze(1) + # When num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps, + # old_per_token_logps == per_token_logps. In this case we can skip its computation + # (see _generate_and_score_completions) and instead use per_token_logps.detach(). + # The exception is when using vLLM, where we always compute old_per_token_logps + # for importance sampling + old_per_token_logps = inputs.get("old_per_token_logps") + old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps + + if self.off_policy_mask_threshold is not None: + # OPSM should use inference-time logprobs to detect both sources of off-policyness: + # 1. Drift from gradient updates (always present) + # 2. Drift from training-inference mismatch (when using vLLM) + # When using vLLM, prioritize sampling_per_token_logps, otherwise use old_per_token_logps + sampling_per_token_logps = inputs.get("sampling_per_token_logps", old_per_token_logps) + + off_policy_mask = self.get_off_policy_mask( + advantages=advantages, + per_token_logps=per_token_logps, + sampling_per_token_logps=sampling_per_token_logps, + mask=mask, + off_policy_threshold=self.off_policy_mask_threshold, + ) + + log_ratio = per_token_logps - old_per_token_logps + if self.importance_sampling_level == "token": + log_importance_weights = log_ratio + elif self.importance_sampling_level == "sequence": + log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0) + log_importance_weights = log_importance_weights.unsqueeze(-1) + else: + raise ValueError( + f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " + "and 'sequence'." + ) + + coef_1 = torch.exp(log_importance_weights) + + # Compute the KL divergence between the model and the reference model + if self.beta != 0.0: + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = ( + torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + ) + # Importance sampling correction for the KL divergence + if self.args.use_bias_correction_kl: + per_token_kl = per_token_kl * coef_1 + + # From here, log_importance_weights (and all subsequent tensors, coef_1, coef_2, etc.) shape depends on + # importance_sampling_level: "token" level: (B, T); "sequence" level: (B, 1) + if self.loss_type == "cispo": + clamped_ratios = torch.clamp(coef_1, max=self.epsilon_high).detach() + per_token_loss = -clamped_ratios * advantages * per_token_logps + elif self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + # Two-sided clipping + if self.args.delta is not None: + coef_1 = torch.clamp(coef_1, max=self.args.delta) + + per_token_loss1 = coef_1 * advantages + per_token_loss2 = coef_2 * advantages + per_token_loss = -torch.min(per_token_loss1, per_token_loss2) + elif self.loss_type == "sapo": + temperatures = torch.where(advantages > 0, self.args.sapo_temperature_pos, self.args.sapo_temperature_neg) + soft_coef_1 = torch.sigmoid(temperatures * (coef_1 - 1)) * 4 / temperatures + per_token_loss = -soft_coef_1 * advantages + elif self.loss_type == "vespo": + phi_seq = self.get_gamma_weights( + advantages=advantages, + log_ratio_per_token=log_ratio, + mask=mask, + importance_sampling_ratio=inputs.get("importance_sampling_ratio"), + k_pos=self.args.vespo_k_pos, + lambda_pos=self.args.vespo_lambda_pos, + k_neg=self.args.vespo_k_neg, + lambda_neg=self.args.vespo_lambda_neg, + ) + per_token_loss = -phi_seq * advantages * per_token_logps + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") + + if self.off_policy_mask_threshold is not None: + per_token_loss = per_token_loss * off_policy_mask + + if entropy_mask is not None: + per_token_loss = per_token_loss * entropy_mask + + if self.use_vllm and self.vllm_importance_sampling_correction and self.loss_type != "vespo": + per_token_loss = per_token_loss * inputs["importance_sampling_ratio"] + + if self.beta != 0.0: + per_token_loss = per_token_loss + self.beta * per_token_kl + + mode = "train" if self.model.training else "eval" + if self.loss_type in ["grpo", "sapo"]: + loss = ((per_token_loss * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() + normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval + loss = loss / normalizer + elif self.loss_type == "bnpo": + loss = (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0) + normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval + loss = loss / normalizer + elif self.loss_type == "dr_grpo": + loss = (per_token_loss * mask).sum() / (per_token_loss.size(0) * self.max_completion_length) + normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval + loss = loss / normalizer + elif self.loss_type in ["cispo", "dapo", "vespo"]: + normalizer = inputs["num_items_in_batch"] / self.accelerator.num_processes + loss = (per_token_loss * mask).sum() / normalizer + elif self.loss_type == "luspo": + # Unless importance_sampling_level="token" (not recommended here), per_token_loss is expected to be (B, 1) + loss = (per_token_loss * mask.sum(1, keepdim=True)).mean() + normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 + loss = loss / normalizer + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") + + # The policy loss above is scaled for gradient accumulation (HF auto-scaling is off here), so scale aux too + if self.aux_loss_enabled: + normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 + loss = loss + self.router_aux_loss_coef * aux_loss / normalizer + self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) + + # Log the metrics + completion_token_count = mask.sum().clamp(min=1.0) + + def masked_batch_mean(x): + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return x.mean() + else: + return (x * mask).sum() / completion_token_count + + if self.beta != 0.0: + mean_kl = masked_batch_mean(per_token_kl) + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + mean_entropy = masked_batch_mean(entropies) + self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) + + if self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) + is_region_clipped = is_low_clipped | is_high_clipped + + low_clip = masked_batch_mean(is_low_clipped.float()) + high_clip = masked_batch_mean(is_high_clipped.float()) + clip_ratio = masked_batch_mean(is_region_clipped.float()) + + gathered_low_clip = self.accelerator.gather(low_clip) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(high_clip) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(clip_ratio) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + elif self.loss_type == "cispo": + is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0) + cispo_clip_ratio = masked_batch_mean(is_cispo_clipped.float()) + gathered_cispo_clip_ratio = self.accelerator.gather(cispo_clip_ratio) + self._metrics[mode]["cispo_clip_ratio"].append(gathered_cispo_clip_ratio.nanmean().item()) + elif self.loss_type == "vespo": + gathered_phi_seq = self.accelerator.gather(phi_seq) + self._metrics[mode]["vespo/phi_seq_mean"].append(gathered_phi_seq.nanmean().item()) + + return loss + + # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and + # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + # Average the metrics + metrics = {} + for key, val in self._metrics[mode].items(): + # Filter out NaN values before averaging. A reward function that returns None for all samples + # in a batch produces NaN for that batch's metric. With logging_steps > 1, a naive sum()/len() + # would let a single NaN contaminate valid data from other batches. Only return None when no + # valid values remain (e.g. JSON loggers crash on float NaN). + valid = [v for v in val if not math.isnan(v)] + metrics[key] = sum(valid) / len(valid) if valid else None + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._logs["prompt"], + self._logs["completion"], + self._logs["rewards"], + self._logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + extra=dict(self._logs["extra"]), + ) + + logging_backends = [] + if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: + logging_backends.append(wandb) + if self.args.report_to and "trackio" in self.args.report_to: + logging_backends.append(trackio) + + table = { + "step": [self.state.global_step] * len(self._logs["prompt"]), + "prompt": self._logs["prompt"], + "completion": self._logs["completion"], + **self._logs["rewards"], + **self._logs["extra"], + "advantage": self._logs["advantages"], + } + + df_base = pd.DataFrame(table) + df_base.to_parquet( + os.path.join( + self.args.output_dir, + "completions", + f"completions_{self.state.global_step:05d}.parquet", + ) + ) + + images_raw = self._logs["images"] or [] + + for logging_backend in logging_backends: + if images_raw: + images = [] + for image_list in self._logs["images"]: + if image_list: + images.append([logging_backend.Image(image) for image in image_list]) + else: + images.append([]) + df = pd.concat( + [df_base, pd.Series(images, name="image")], + axis=1, + copy=False, + ) + else: + df = df_base + + if self.log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + + logging_backend.log({"completions": logging_backend.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) +class UnslothGRPOTrainer(_UnslothGRPOTrainer): + """ + + Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the + paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language + Models](https://huggingface.co/papers/2402.03300). + + Example: + + ```python + >>> from trl import GRPOTrainer + >>> from trl.rewards import accuracy_reward + >>> from datasets import load_dataset + + >>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train") + + >>> trainer = GRPOTrainer( + ... model="Qwen/Qwen2.5-0.5B-Instruct", + ... reward_funcs=accuracy_reward, + ... train_dataset=dataset, + ... ) + >>> trainer.train() + ``` + + Args: + model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `.from_pretrained` (where `` is derived from the model + config) with the keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + - A [`~peft.PeftModel`] object. Only causal language models are supported. + reward_funcs (`RewardFunc | list[RewardFunc]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can be either synchronous or asynchronous and can also return `None` when the reward is + not applicable to those samples. This is useful for multi-task training where different reward + functions apply to different types of samples. When a reward function returns `None` for a sample, + that reward function is excluded from the reward calculation for that sample. For more details, see + [Using a custom reward + function](#using-a-custom-reward-function). + + The trainer's state is also passed to the reward function. The trainer's state is an instance of + [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the + reward function's signature. + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + args ([`GRPOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + tools (list of `Callable`, *optional*): + A list of callable tool functions (sync or async) that the model can invoke during generation. Each tool + should be a standard Python function with properly type-hinted arguments and return values, and a + Google-style docstring describing its purpose, arguments, and return value. For more details, see: + https://huggingface.co/docs/transformers/en/chat_extras#passing-tools. The model uses the function's name, + type hints, and docstring to determine how to call it. Ensure that the model's chat template supports tool + use and that it has been fine-tuned for tool calling. + rollout_func (`RolloutFunc`, *optional*): + Function to use for generating completions. It receives the list of prompts allocated to the current + process and the trainer instance. It must return a dict with `"prompt_ids"`, `"completion_ids"`, and + `"logprobs"` fields, and can optionally return `"logprob_token_ids"` (same shape as `"logprobs"`). Any + other fields are forwarded to the reward functions. The function receives the raw per-process prompt slice + with no duplication; it is responsible for returning the correct number of completions per prompt (see + `num_generations` / `num_generations_eval` on the trainer). This feature is experimental and may change or + be removed at any time without prior notice. + environment_factory (`EnvironmentFactory`, *optional*): + A callable that creates and returns an environment instance. The environment class should define methods + that can be invoked as tools during generation. Each method should comply with the same requirements as the + `tools` described above. If `environment_factory` is provided, an instance of the environment is created + for each generation in the batch, allowing for parallel and independent interactions. The environment must + also implement a callable `reset` method that can be used to reset state between generations. The `reset` + method should return either `None` or a string: when it returns a string, that string is appended to the + last user message before generation. This feature is experimental and may change or be removed at any time + without prior notice. + + """ + def __init__( + self, + model, + reward_funcs, + args = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + reward_processing_classes = None, + callbacks = None, + peft_config = None, + tools = None, + rollout_func = None, + environment_factory = None, + **kwargs + ): + if args is None: args = UnslothGRPOConfig() + use_bf16 = getattr(args, 'bf16', False) + if type(use_bf16) is not bool: use_bf16 = False + use_fp16 = getattr(args, 'fp16', False) + if type(use_fp16) is not bool: use_fp16 = False + force_float32 = False + full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' + if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): + print('Unsloth: Switching to float32 training since model cannot work with float16') + force_float32 = True + mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') + dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) + if dtype is None: dtype = model.get_input_embeddings().weight.dtype + from unsloth_zoo.utils import _get_dtype + dtype = _get_dtype(dtype) + float16 = dtype == torch.float16 + if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') + if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') + if force_float32: + # Forced float32 training + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': + # Mixed precision training + args.fp16 = float16 + args.bf16 = not float16 + os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' + # args.mixed_precision is a new argument which needs to be set now + elif mixed_precision_dtype == 'bfloat16': + # Both False since bfloat16 full finetuning doesn't do any autocasting. + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + + if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': + args.eval_strategy = 'steps' + if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 + ga_steps = getattr(args, 'gradient_accumulation_steps', None) + if ga_steps is not None and ga_steps > 1: + from transformers import __version__ as transformers_version + if Version(transformers_version) <= Version('4.45.2'): + print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' + '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') + if getattr(args, 'eval_strategy', 'no') != 'no': + eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) + if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size + if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps + fp16_full_eval = getattr(args, 'fp16_full_eval', False) + if type(fp16_full_eval) is not bool: fp16_full_eval = False + bf16_full_eval = getattr(args, 'bf16_full_eval', False) + if type(bf16_full_eval) is not bool: bf16_full_eval = False + if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True + if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False + if force_float32: + args.bf16_full_eval = False + args.fp16_full_eval = False + elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': + args.bf16_full_eval = True + args.fp16_full_eval = False + elif not bf16_full_eval and not fp16_full_eval: + args.bf16_full_eval = args.bf16 + args.fp16_full_eval = args.fp16 + _output_logits = False + if locals().get('compute_metrics', None) is not None: _output_logits = True + if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True + if _output_logits: + os.environ['UNSLOTH_RETURN_LOGITS'] = '1' + if model is not None: + _warnings_issued = getattr(model, 'warnings_issued', None) + if _warnings_issued is None: + model.warnings_issued = {} + elif not isinstance(_warnings_issued, dict): + try: + model.warnings_issued = dict(_warnings_issued) + except Exception: + model.warnings_issued = {} + if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): + pass + else: + model_max_seq_length = getattr(model, 'max_seq_length', None) + args_max_seq_length = getattr(args, 'max_seq_length', None) + if args_max_seq_length is None and model_max_seq_length is not None: + max_seq_length = model.max_seq_length + if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length + elif args_max_seq_length is not None and model_max_seq_length is not None: + if args_max_seq_length > model_max_seq_length: + print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' + 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') + args.max_seq_length = model_max_seq_length + if model is not None and hasattr(model, 'for_training'): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' + if 'processing_class' in locals(): + if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' + if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' + other_metrics = [] + if not isinstance(reward_funcs, list): _reward_funcs = [reward_funcs] + else: _reward_funcs = reward_funcs + for reward_func in _reward_funcs: + try: + reward_func_name = reward_func.__name__ + if True: + other_metrics.append(f'rewards/{reward_func_name}/mean') + if True: + other_metrics.append(f'rewards/{reward_func_name}/std') + if False: + other_metrics.append(f'rewards/{reward_func_name}') + except: pass + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('grpo_trainer', other_metrics) + + # [TODO] Fix up DataParallel multiplying batch sizes + # [TODO] DDP works, but DP seems to not work? [TODO] + if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: + if getattr(args, "_n_gpu", 1) != 1: + args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + super().__init__( + model = model, + reward_funcs = reward_funcs, + args = args, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + reward_processing_classes = reward_processing_classes, + callbacks = callbacks, + peft_config = peft_config, + tools = tools, + rollout_func = rollout_func, + environment_factory = environment_factory,**kwargs) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() + if hasattr(self, 'neftune_hook_handle'): + self.neftune_hook_handle.remove() + if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle + if getattr(args, 'neftune_noise_alpha', None) is not None: + model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha + pass + if hasattr(self, 'accelerator'): + scaler = self.accelerator.scaler + current_model = model + while hasattr(current_model, 'model'): + current_model.accelerator_scaler = scaler + current_model = current_model.model + current_model.accelerator_scaler = scaler + pass + if hasattr(self, 'train'): + self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) + pass + if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): + _vllm_tok = self.llm.get_tokenizer() + _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) + if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: + _vllm_tok.chat_template = _pc.chat_template + pass + +pass + + +if hasattr(logger, "addFilter"): + import logging + class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not (self.text in x.getMessage()) + pass + logger.addFilter(HideLoggingMessage("`use_cache=True`")) + diff --git a/unsloth_compiled_cache/UnslothKTOTrainer.py b/unsloth_compiled_cache/UnslothKTOTrainer.py new file mode 100644 index 0000000000000000000000000000000000000000..66722ea4fd499a560d551075caa5e23bf740385a --- /dev/null +++ b/unsloth_compiled_cache/UnslothKTOTrainer.py @@ -0,0 +1,2048 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.experimental.kto.kto_trainer import (Any, AutoProcessor, Callable, DataCollator, DataCollatorForUnpairedPreference, DataCollatorForVisionUnpairedPreference, DataLoader, Dataset, EvalLoopOutput, F, Hasher, IterableDataset, IterableDatasetDict, KTOConfig, KTOTrainer, LoraConfig, PartialState, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Sampler, SequentialSampler, SyncRefModelCallback, TrainerCallback, Version, _BaseTrainer, _get_kl_completion_ids, apply_chat_template, concatenate_datasets, contextlib, create_model_from_path, dataclass, defaultdict, disable_dropout_in_model, disable_gradient_checkpointing, extract_prompt, flush_left, get_act_offloading_ctx_manager, get_config_model_id, get_dataset_column_names, get_peft_model, has_length, hash_module, is_conversational, is_liger_kernel_available, is_peft_available, is_peft_model, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, selective_log_softmax, textwrap, torch, tqdm, transformers, unpair_preference_dataset, use_adapter, AutoProcessor, Callable, DataCollator, DataCollatorForUnpairedPreference, DataCollatorForVisionUnpairedPreference, Dataset, EvalLoopOutput, F, IterableDataset, IterableDatasetDict, KTOConfig, KTOTrainer, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SyncRefModelCallback, TrainerCallback, Version, contextlib, create_model_from_path, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, is_liger_kernel_available, is_peft_available, is_peft_model, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, torch, transformers, unpair_preference_dataset, F, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch) + + +import os +import math +import logging +from typing import * +from dataclasses import dataclass, field +from packaging.version import Version +import torch +import numpy as np +from contextlib import nullcontext +from torch.nn import functional as F +import inspect +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling +from transformers.training_args import ParallelMode +from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize + +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass +# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain +# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. +try: + from unsloth.models._utils import _unsloth_reset_stray_compile_cache +except Exception: + def _unsloth_reset_stray_compile_cache(self): pass +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Drop any torch.compile graph cache poisoned by a stray pre-train forward. + try: + _unsloth_reset_stray_compile_cache(self) + except Exception: + pass + # Finish the previous W&B run if this is a subsequent train() call. + # We do this at the START of train() (not the end) so that + # evaluate() / log() still work after train() completes. + # HF's WandbCallback.setup() will call wandb.init() for the new run. + # See: https://github.com/unslothai/unsloth/issues/3954 + if getattr(self, '_unsloth_training_completed', False): + try: + import wandb + if wandb.run is not None: + wandb.finish() + # Reset HF's WandbCallback so it calls wandb.init() for the new run + for cb in self.callback_handler.callbacks: + if type(cb).__name__ == 'WandbCallback': + cb._initialized = False + break + except: + pass + # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + output = f(self, *args, **kwargs) + # Restore previous mode when possible + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass + # Mark that training completed so the next train() call can + # finish this W&B run before starting a new one + self._unsloth_training_completed = True + return output + return wrapper +pass + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_hidden_states_selective_log_softmax( + hidden_states: torch.Tensor, + lm_head: torch.Tensor, + index: torch.Tensor, + chunks: int = 4, + logit_scale_multiply: float = 0.0, + logit_scale_divide: float = 0.0, + logit_softcapping: float = 0.0, + temperature: float = 1.0, +) -> torch.Tensor: + # All Unsloth Zoo code licensed under AGPL3 + flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + flat_index = index.reshape(-1) + + chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) + chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) + + all_per_token_logps = [] + + for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): + chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() + + if logit_scale_multiply != 0.0: + chunk_logits = chunk_logits * logit_scale_multiply + if logit_scale_divide != 0.0: + chunk_logits = chunk_logits / logit_scale_divide + if logit_softcapping != 0.0: + chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) + + chunk_logits = chunk_logits.to(torch.float32) + + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + + selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + + all_per_token_logps = torch.concat(all_per_token_logps) + + all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) + return all_per_token_logps + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_selective_log_softmax( + logits, + index, + temperature: float = 1.0, + chunks: int = 4, +): + chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) + chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) + all_per_token_logps = [] + # Per-chunk selective_log_softmax. + for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): + chunk_logits = chunk_logits.to(torch.float32) + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + pass + all_per_token_logps = torch.concat(all_per_token_logps) + all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) + return all_per_token_logps + +def calculate_pad_tokens_in_prompt( + input_ids: torch.Tensor, + logits_to_keep: int, + pad_token_id: int +) -> torch.Tensor: + """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" + if logits_to_keep >= input_ids.shape[1]: + raise ValueError("logits_to_keep must be smaller than the sequence length.") + + prompt_section = input_ids[:, :-logits_to_keep] + + padding_mask = (prompt_section == pad_token_id) + + pad_token_counts = padding_mask.sum(dim=1) + + return pad_token_counts + +def create_completion_attention_mask( + completion_input_ids: torch.Tensor, + left_pad_tokens_per_prompt: torch.Tensor, + max_left_pad: int, + pad_token_id: int +) -> torch.Tensor: + """Build a completion mask that zeros leading prompt and trailing pad tokens. + + For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) + this returns [0,0,0,1,1,1,0,0,0]. + """ + batch_size, completion_len = completion_input_ids.shape + device = completion_input_ids.device + + num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt + + indices = torch.arange(completion_len, device=device).unsqueeze(0) + shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) + + non_padding_mask = (completion_input_ids != pad_token_id) + + final_mask = shift_mask & non_padding_mask + + return final_mask + +def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: + """Move all padding tokens in each sequence to the right.""" + mask = (tensor != pad_id) + # stable=True since the binary mask is unordered. + sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) + packed_tensor = torch.gather(tensor, 1, sorted_indices) + return packed_tensor + +def align_logprobs_with_mask( + logprob_tensor: torch.Tensor, + attention_mask: torch.Tensor, + pad_value: float = 0.0 +) -> torch.Tensor: + """Align a log probability tensor with a given attention mask.""" + + device = logprob_tensor.device + batch_size, logprob_seq_len = logprob_tensor.shape + mask_seq_len = attention_mask.shape[1] + + padded_logprobs = torch.full( + attention_mask.shape, + fill_value=pad_value, + dtype=logprob_tensor.dtype, + device=device + ) + + left_pad_counts = torch.argmax(attention_mask, dim=1) + + cols = torch.arange(logprob_seq_len, device=device) + dest_indices = left_pad_counts.unsqueeze(1) + cols + + # Destination row indices, shape [batch_size, logprob_seq_len]. + row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) + + # Keep only in-bounds destinations, then scatter via advanced indexing. + valid_mask = dest_indices < mask_seq_len + valid_rows = row_indices[valid_mask] + valid_cols = dest_indices[valid_mask] + valid_vals = logprob_tensor[valid_mask] + padded_logprobs[valid_rows, valid_cols] = valid_vals + + return padded_logprobs + +def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" + if tool_mask is None: + return completion_mask + if tool_mask.shape[0] != completion_mask.shape[0]: + raise ValueError("tool_mask batch size must match completion_mask batch size.") + + tool_mask = tool_mask.to(device=completion_mask.device) + if tool_mask.shape == completion_mask.shape: + aligned_tool_mask = tool_mask + else: + aligned_tool_mask = align_logprobs_with_mask( + tool_mask, + completion_mask, + pad_value=0, + ) + return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) + +def autotune_batch_and_chunks( + total_input_rows, + seq_len, + hidden_size, + vocab_size, + dtype_bytes=16, + multiplier=None +): + if multiplier is None: + final_m = max(4, seq_len // 4096) + else: + final_m = multiplier + + if torch.cuda.is_available(): + free_bytes, _ = torch.cuda.mem_get_info() + limit_gb = (free_bytes / (1024**3))*.80 + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + # XPU: estimate free memory as total - reserved. + total_mem = torch.xpu.get_device_properties(0).total_memory + reserved_mem = torch.xpu.memory_reserved() + free_bytes = total_mem - reserved_mem + limit_gb = (free_bytes / (1024**3)) * 0.80 + else: + # Fallback: assume 8GB available. + limit_gb = 8.0 + + bytes_to_gb = 1024**3 + + b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) + + hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb + + base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb + logits_gb = base_logits / final_m + + total_mem_gb = hidden_gb + logits_gb + + valid_mask = total_mem_gb <= limit_gb + valid_indices = torch.nonzero(valid_mask, as_tuple=False) + + if valid_indices.shape[0] == 0: + #This means your GPU will OOM + return 4, final_m + + best_idx = valid_indices[0].item() + final_b = int(b_vals[best_idx].item()) + + return final_b, final_m + +def sanitize_logprob(logprob): + """Local port of trl.scripts.vllm_serve.sanitize_logprob. + Filters NaN logprobs from vLLM outputs.""" + value = logprob.logprob + if math.isnan(value): + logging.getLogger(__name__).warning( + f"Generated NaN logprob, token logprob '{logprob}' will be ignored" + ) + return None + return value +@dataclass +class UnslothKTOConfig(KTOConfig): + """ + KTOConfig(output_dir: str | None = None, per_device_train_batch_size: int = 8, num_train_epochs: float = 3.0, max_steps: int = -1, learning_rate: float = 1e-06, lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear', lr_scheduler_kwargs: dict | str | None = None, warmup_steps: float = 0, optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused', optim_args: str | None = None, weight_decay: float = 0.0, adam_beta1: float = 0.9, adam_beta2: float = 0.999, adam_epsilon: float = 1e-08, optim_target_modules: None | str | list[str] = None, gradient_accumulation_steps: int = 1, average_tokens_across_devices: bool = True, max_grad_norm: float = 1.0, label_smoothing_factor: float = 0.0, bf16: bool | None = None, fp16: bool = False, bf16_full_eval: bool = False, fp16_full_eval: bool = False, tf32: bool | None = None, gradient_checkpointing: bool = True, gradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = None, torch_compile: bool = False, torch_compile_backend: str | None = None, torch_compile_mode: str | None = None, use_liger_kernel: bool = False, liger_kernel_config: dict[str, bool] | None = None, use_cache: bool = False, neftune_noise_alpha: float | None = None, torch_empty_cache_steps: int | None = None, auto_find_batch_size: bool = False, logging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps', logging_steps: float = 10, logging_first_step: bool = False, log_on_each_node: bool = True, logging_nan_inf_filter: bool = True, include_num_input_tokens_seen: str | bool = 'no', log_level: str = 'passive', log_level_replica: str = 'warning', disable_tqdm: bool | None = None, report_to: None | str | list[str] = 'none', run_name: str | None = None, project: str = 'huggingface', trackio_space_id: str | None = 'trackio', eval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no', eval_steps: float | None = None, eval_delay: float = 0, per_device_eval_batch_size: int = 8, prediction_loss_only: bool = False, eval_on_start: bool = False, eval_do_concat_batches: bool = True, eval_use_gather_object: bool = False, eval_accumulation_steps: int | None = None, include_for_metrics: list[str] = , batch_eval_metrics: bool = False, save_only_model: bool = False, save_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps', save_steps: float = 500, save_on_each_node: bool = False, save_total_limit: int | None = None, enable_jit_checkpoint: bool = False, push_to_hub: bool = False, hub_token: str | None = None, hub_private_repo: bool | None = None, hub_model_id: str | None = None, hub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save', hub_always_push: bool = False, hub_revision: str | None = None, load_best_model_at_end: bool = False, metric_for_best_model: str | None = None, greater_is_better: bool | None = None, ignore_data_skip: bool = False, restore_callback_states_from_checkpoint: bool = False, full_determinism: bool = False, seed: int = 42, data_seed: int | None = None, use_cpu: bool = False, accelerator_config: dict | str | None = None, parallelism_config: accelerate.parallelism_config.ParallelismConfig | None = None, dataloader_drop_last: bool = False, dataloader_num_workers: int = 0, dataloader_pin_memory: bool = True, dataloader_persistent_workers: bool = False, dataloader_prefetch_factor: int | None = None, remove_unused_columns: bool = True, label_names: list[str] | None = None, train_sampling_strategy: str = 'sequential', length_column_name: str = 'length', ddp_find_unused_parameters: bool | None = None, ddp_bucket_cap_mb: int | None = None, ddp_broadcast_buffers: bool | None = None, ddp_backend: str | None = None, ddp_timeout: int = 1800, fsdp: list[transformers.trainer_utils.FSDPOption] | str | None = None, fsdp_config: dict[str, typing.Any] | str | None = None, deepspeed: dict | str | None = None, debug: str | list[transformers.debug_utils.DebugOption] = '', skip_memory_metrics: bool = True, do_train: bool = False, do_eval: bool = False, do_predict: bool = False, resume_from_checkpoint: str | None = None, warmup_ratio: float | None = None, logging_dir: str | None = None, local_rank: int = -1, model_init_kwargs: dict[str, typing.Any] | str | None = None, trust_remote_code: bool = False, disable_dropout: bool = True, dataset_num_proc: int | None = None, max_length: int | None = 1024, pad_to_multiple_of: int | None = None, precompute_ref_log_probs: bool = False, precompute_ref_batch_size: int | None = None, loss_type: str = 'kto', beta: float = 0.1, desirable_weight: float = 1.0, undesirable_weight: float = 1.0, activation_offloading: bool = False, sync_ref_model: bool = False, ref_model_mixup_alpha: float = 0.6, ref_model_sync_steps: int = 512) + """ + vllm_sampling_params: Optional[Any] = field( + default = None, + metadata = {'help': 'vLLM SamplingParams'}, + ) + unsloth_num_chunks : Optional[int] = field( + default = -1, + metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, + ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {'help': 'Multiplier for chunked logit computations.'}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, + ) + max_seq_length : Optional[int] = field( + default = None, + metadata = {'help': 'Maximum sequence length to truncate to.'}, + ) + def __init__( + self, + output_dir = None, + per_device_train_batch_size = 4, + num_train_epochs = 3.0, + max_steps = -1, + learning_rate = 5e-05, + lr_scheduler_type = 'linear', + lr_scheduler_kwargs = None, + warmup_steps = 0.1, + optim = 'adamw_8bit', + optim_args = None, + weight_decay = 0.001, + adam_beta1 = 0.9, + adam_beta2 = 0.999, + adam_epsilon = 1e-08, + optim_target_modules = None, + gradient_accumulation_steps = 2, + average_tokens_across_devices = True, + max_grad_norm = 1.0, + label_smoothing_factor = 0.0, + bf16 = False, + fp16 = False, + bf16_full_eval = False, + fp16_full_eval = False, + tf32 = None, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = None, + torch_compile = False, + torch_compile_backend = None, + torch_compile_mode = None, + use_liger_kernel = False, + liger_kernel_config = None, + use_cache = False, + neftune_noise_alpha = None, + torch_empty_cache_steps = 250, + auto_find_batch_size = False, + logging_strategy = 'steps', + logging_steps = 1, + logging_first_step = False, + log_on_each_node = True, + logging_nan_inf_filter = False, + include_num_input_tokens_seen = False, + log_level = 'passive', + log_level_replica = 'warning', + disable_tqdm = None, + report_to = 'none', + run_name = None, + project = 'huggingface', + trackio_space_id = 'trackio', + eval_strategy = 'no', + eval_steps = None, + eval_delay = 0, + per_device_eval_batch_size = 4, + prediction_loss_only = False, + eval_on_start = False, + eval_do_concat_batches = True, + eval_use_gather_object = False, + eval_accumulation_steps = 2, + batch_eval_metrics = False, + save_only_model = False, + save_strategy = 'steps', + save_steps = 500, + save_on_each_node = False, + save_total_limit = None, + enable_jit_checkpoint = False, + push_to_hub = False, + hub_token = None, + hub_private_repo = None, + hub_model_id = None, + hub_strategy = 'every_save', + hub_always_push = False, + hub_revision = None, + load_best_model_at_end = False, + metric_for_best_model = None, + greater_is_better = None, + ignore_data_skip = False, + restore_callback_states_from_checkpoint = False, + full_determinism = False, + seed = 3407, + data_seed = 3407, + use_cpu = False, + accelerator_config = None, + parallelism_config = None, + dataloader_drop_last = False, + dataloader_num_workers = 0, + dataloader_pin_memory = True, + dataloader_persistent_workers = False, + dataloader_prefetch_factor = None, + remove_unused_columns = True, + label_names = None, + train_sampling_strategy = 'sequential', + length_column_name = 'length', + ddp_find_unused_parameters = None, + ddp_bucket_cap_mb = None, + ddp_broadcast_buffers = None, + ddp_backend = None, + ddp_timeout = 1800, + fsdp = None, + fsdp_config = None, + deepspeed = None, + debug = '', + skip_memory_metrics = True, + do_train = False, + do_eval = False, + do_predict = False, + resume_from_checkpoint = None, + warmup_ratio = None, + logging_dir = None, + local_rank = -1, + model_init_kwargs = None, + trust_remote_code = False, + disable_dropout = True, + dataset_num_proc = None, + max_length = 1024, + pad_to_multiple_of = None, + precompute_ref_log_probs = False, + precompute_ref_batch_size = None, + loss_type = 'kto', + beta = 0.1, + desirable_weight = 1.0, + undesirable_weight = 1.0, + activation_offloading = False, + sync_ref_model = False, + ref_model_mixup_alpha = 0.6, + ref_model_sync_steps = 512, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, + max_seq_length = None, + **kwargs, + ): + if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') + if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') + if num_train_epochs is None: + num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override + if output_dir is None and save_strategy == 'steps' and save_steps == 500: + output_dir = 'unsloth_training_checkpoints' + save_strategy = 'no' + import multiprocessing as _mp + if dataset_num_proc is None: + if _mp.get_start_method() != 'fork': + dataset_num_proc = None + else: + import psutil + dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64) + memory_gb_left = psutil.virtual_memory().available / (1024**3) + if memory_gb_left <= 2: dataset_num_proc = 1 + else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left)) + if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': + from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION + if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: + from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE + pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE + + + super().__init__( + output_dir = output_dir, + per_device_train_batch_size = per_device_train_batch_size, + num_train_epochs = num_train_epochs, + max_steps = max_steps, + learning_rate = learning_rate, + lr_scheduler_type = lr_scheduler_type, + lr_scheduler_kwargs = lr_scheduler_kwargs, + warmup_steps = warmup_steps, + optim = optim, + optim_args = optim_args, + weight_decay = weight_decay, + adam_beta1 = adam_beta1, + adam_beta2 = adam_beta2, + adam_epsilon = adam_epsilon, + optim_target_modules = optim_target_modules, + gradient_accumulation_steps = gradient_accumulation_steps, + average_tokens_across_devices = average_tokens_across_devices, + max_grad_norm = max_grad_norm, + label_smoothing_factor = label_smoothing_factor, + bf16 = bf16, + fp16 = fp16, + bf16_full_eval = bf16_full_eval, + fp16_full_eval = fp16_full_eval, + tf32 = tf32, + gradient_checkpointing = gradient_checkpointing, + gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, + torch_compile = torch_compile, + torch_compile_backend = torch_compile_backend, + torch_compile_mode = torch_compile_mode, + use_liger_kernel = use_liger_kernel, + liger_kernel_config = liger_kernel_config, + use_cache = use_cache, + neftune_noise_alpha = neftune_noise_alpha, + torch_empty_cache_steps = torch_empty_cache_steps, + auto_find_batch_size = auto_find_batch_size, + logging_strategy = logging_strategy, + logging_steps = logging_steps, + logging_first_step = logging_first_step, + log_on_each_node = log_on_each_node, + logging_nan_inf_filter = logging_nan_inf_filter, + include_num_input_tokens_seen = include_num_input_tokens_seen, + log_level = log_level, + log_level_replica = log_level_replica, + disable_tqdm = disable_tqdm, + report_to = report_to, + run_name = run_name, + project = project, + trackio_space_id = trackio_space_id, + eval_strategy = eval_strategy, + eval_steps = eval_steps, + eval_delay = eval_delay, + per_device_eval_batch_size = per_device_eval_batch_size, + prediction_loss_only = prediction_loss_only, + eval_on_start = eval_on_start, + eval_do_concat_batches = eval_do_concat_batches, + eval_use_gather_object = eval_use_gather_object, + eval_accumulation_steps = eval_accumulation_steps, + batch_eval_metrics = batch_eval_metrics, + save_only_model = save_only_model, + save_strategy = save_strategy, + save_steps = save_steps, + save_on_each_node = save_on_each_node, + save_total_limit = save_total_limit, + enable_jit_checkpoint = enable_jit_checkpoint, + push_to_hub = push_to_hub, + hub_token = hub_token, + hub_private_repo = hub_private_repo, + hub_model_id = hub_model_id, + hub_strategy = hub_strategy, + hub_always_push = hub_always_push, + hub_revision = hub_revision, + load_best_model_at_end = load_best_model_at_end, + metric_for_best_model = metric_for_best_model, + greater_is_better = greater_is_better, + ignore_data_skip = ignore_data_skip, + restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, + full_determinism = full_determinism, + seed = seed, + data_seed = data_seed, + use_cpu = use_cpu, + accelerator_config = accelerator_config, + parallelism_config = parallelism_config, + dataloader_drop_last = dataloader_drop_last, + dataloader_num_workers = dataloader_num_workers, + dataloader_pin_memory = dataloader_pin_memory, + dataloader_persistent_workers = dataloader_persistent_workers, + dataloader_prefetch_factor = dataloader_prefetch_factor, + remove_unused_columns = remove_unused_columns, + label_names = label_names, + train_sampling_strategy = train_sampling_strategy, + length_column_name = length_column_name, + ddp_find_unused_parameters = ddp_find_unused_parameters, + ddp_bucket_cap_mb = ddp_bucket_cap_mb, + ddp_broadcast_buffers = ddp_broadcast_buffers, + ddp_backend = ddp_backend, + ddp_timeout = ddp_timeout, + fsdp = fsdp, + fsdp_config = fsdp_config, + deepspeed = deepspeed, + debug = debug, + skip_memory_metrics = skip_memory_metrics, + do_train = do_train, + do_eval = do_eval, + do_predict = do_predict, + resume_from_checkpoint = resume_from_checkpoint, + warmup_ratio = warmup_ratio, + logging_dir = logging_dir, + local_rank = local_rank, + model_init_kwargs = model_init_kwargs, + trust_remote_code = trust_remote_code, + disable_dropout = disable_dropout, + dataset_num_proc = dataset_num_proc, + max_length = max_length, + pad_to_multiple_of = pad_to_multiple_of, + precompute_ref_log_probs = precompute_ref_log_probs, + precompute_ref_batch_size = precompute_ref_batch_size, + loss_type = loss_type, + beta = beta, + desirable_weight = desirable_weight, + undesirable_weight = undesirable_weight, + activation_offloading = activation_offloading, + sync_ref_model = sync_ref_model, + ref_model_mixup_alpha = ref_model_mixup_alpha, + ref_model_sync_steps = ref_model_sync_steps,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier + self.max_seq_length = max_seq_length + # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+ + if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: + if 'use_reentrant' in self.gradient_checkpointing_kwargs: + del self.gradient_checkpointing_kwargs['use_reentrant'] + +pass + +class _UnslothKTOTrainer(_BaseTrainer): + """ + Initialize KTOTrainer. + + Args: + model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `.from_pretrained` (where `` is derived from the model + config) with the keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + - A [`~peft.PeftModel`] object. Only causal language models are supported. + ref_model ([`~transformers.PreTrainedModel`], *optional*): + Reference model used to compute the reference log probabilities. + + - If provided, this model is used directly as the reference policy. + - If `None`, the trainer will automatically use the initial policy corresponding to `model`, i.e. the model + state before KTO training starts. + args ([`experimental.kto.KTOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + data_collator ([`~transformers.DataCollator`], *optional*): + The data collator to use for training. If None is specified, the default data collator + ([`~experimental.kto.kto_trainer.DataCollatorForUnpairedPreference`]) will be used which will pad the + sequences to the maximum length of the sequences in the batch. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + """ + + _tag_names = ["trl", "kto"] + _name = "KTO" + _paper = { + "title": "KTO: Model Alignment as Prospect Theoretic Optimization", + "id": "2402.01306", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{ethayarajh2024kto, + title = {{KTO: Model Alignment as Prospect Theoretic Optimization}}, + author = {Kawin Ethayarajh and Winnie Xu and Niklas Muennighoff and Dan Jurafsky and Douwe Kiela}, + year = 2024, + eprint = {arXiv:2402.01306}, + }"""), + } + + def __init__( + self, + model: "str | PreTrainedModel | PeftModel", + ref_model: PreTrainedModel | None = None, + args: KTOConfig | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + data_collator: DataCollator | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + peft_config: "PeftConfig | None" = None, + compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = KTOConfig(f"{model_name}-KTO") + + if train_dataset is None: + raise ValueError("`train_dataset` is required") + elif isinstance(train_dataset, IterableDataset): + # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate + # batches from multiple processes, leading to mismatch errors. + if args.accelerator_config.dispatch_batches is True: + logger.warning( + "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " + "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " + "`dispatch_batches` in `KTOConfig` or set it to `False`." + ) + args.accelerator_config.dispatch_batches = False + + # Model + if isinstance(model, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + model = create_model_from_path(model, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the KTOConfig, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do + _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. In most cases you should omit `ref_model` and " + "we'll initialize it to a copy of `model` for you." + ) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained( + get_config_model_id(model.config), trust_remote_code=args.trust_remote_code + ) + if isinstance(processing_class, ProcessorMixin): + self._tokenizer = processing_class.tokenizer + self._is_vlm = True + elif isinstance(processing_class, PreTrainedTokenizerBase): + self._tokenizer = processing_class + self._is_vlm = False + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + if self._tokenizer.pad_token is None: + self._tokenizer.pad_token = self._tokenizer.eos_token + + # PEFT + if False: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if is_peft_model(model): + raise ValueError( + "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " + "and unload the existing adapter, save the resulting base model, and then pass that base model along " + "with the new `peft_config` to the trainer." + ) + # Create PEFT model + # ZeRO-3 + PEFT for non-quantized models: + # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. + # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, + # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. + # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. + # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for + # non-quantized bf16 training. + # - See: + # - TRL issue: https://github.com/huggingface/trl/issues/6089 + # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 + # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg + get_peft_model_kwargs = {} + if ( + args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and not _is_quantized_model + and Version(peft.__version__) >= Version("0.12.0") + ): + get_peft_model_kwargs["autocast_adapter_dtype"] = False + model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + + elif is_peft_model(model) and ref_model is None: + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during KTO training. PEFT only + # supports one adapter per model when the LoRA config uses `target_parameters` [see peft#3340], so in that + # case we skip the "ref" adapter and compute the reference log probs with adapters disabled, i.e. with the + # base model. + default_config = model.peft_config["default"] + if isinstance(default_config, LoraConfig) and default_config.target_parameters: + logger.warning( + "PEFT can't add a frozen reference adapter alongside one that uses `target_parameters` " + "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " + "If you wrapped the model only to apply LoRA, pass a `peft_config` to the trainer instead; if you " + "wrapped it deliberately (pretrained adapter or custom init), note that the base model matches " + "your adapter only when it's freshly zero-initialized. If it is, this warning is safe to ignore." + ) + else: + model.add_adapter("ref", default_config) + for name, param in model.named_parameters(): + if ".default." in name: + ref_name = name.replace(".default.", ".ref.") + ref_param = model.get_parameter(ref_name) + ref_param.data.copy_(param.data) + + # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally + # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 + if is_peft_model(model) and args.gradient_checkpointing: + model.enable_input_require_grads() + + # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the + # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by + # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for + # quantized models. See: https://github.com/huggingface/peft/issues/2889 + if _is_quantized_model: + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.bfloat16) + + # Vision dataset detection + dataset_sample = next(iter(train_dataset)) + self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample + if self._is_vision_dataset and not self._is_vlm: + raise ValueError( + "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " + "model does not seem to be a vision-language model. Please check your model and dataset." + ) + if self._is_vision_dataset and args.precompute_ref_log_probs: + raise ValueError( + "`precompute_ref_log_probs=True` is not supported for vision datasets. For vision-language " + "models, all data processing is performed on the fly rather than upfront. " + "Set `precompute_ref_log_probs=False`." + ) + if self._is_vision_dataset and ("chosen" in dataset_sample or "rejected" in dataset_sample): + raise ValueError( + "Vision datasets must be in unpaired format with `completion` and `label` columns. " + "Paired format (`chosen`/`rejected`) is not supported for vision datasets because " + "iterating over the full dataset to unpair it would be too expensive for large image " + "collections. Unpair your dataset first: `dataset = unpair_preference_dataset(dataset)`." + ) + + # Data collator + calculate_kl = args.loss_type not in ["apo_zero_unpaired"] + if data_collator is None and not self._is_vision_dataset: + data_collator = DataCollatorForUnpairedPreference( + pad_token_id=self._tokenizer.pad_token_id, + max_length=args.max_length, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + elif data_collator is None and self._is_vision_dataset: + data_collator = DataCollatorForVisionUnpairedPreference( + processor=processing_class, + max_length=args.max_length, + calculate_kl=calculate_kl, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + + # Training arguments + self.beta = args.beta + self.precompute_ref_logps = args.precompute_ref_log_probs + self.loss_type = args.loss_type + self.desirable_weight = args.desirable_weight + self.undesirable_weight = args.undesirable_weight + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + self.calculate_KL = calculate_kl + if self.calculate_KL and args.train_sampling_strategy != "sequential": + raise ValueError( + f"Loss type `'{args.loss_type}'` estimates the KL divergence term and requires " + f"`train_sampling_strategy='sequential'` because the KL completion for each example is precomputed " + f"against its neighbors in a fixed-order batch; any other strategy breaks that pairing. " + f"Got `train_sampling_strategy='{args.train_sampling_strategy}'`." + ) + if self.calculate_KL and args.per_device_train_batch_size <= 1: + raise ValueError( + "Actual (not effective) batch size must be > 1. KTO will not work properly because the KL term will be equivalent to the implied reward." + ) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + logger.warning( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + ) + + # Dataset + # Skip dataset preparation for VLMs: tokenization and image processing happen on-the-fly in the collator. + if not self._is_vision_dataset: + train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + ) + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # Reference model + if ref_model is None: + if is_peft_model(self.model) or args.precompute_ref_log_probs: + # If PEFT is used, the reference model is not needed since the adapter can be disabled to revert to the + # initial model. If precompute_ref_log_probs is True, the reference model does not need to be kept in + # memory during training. + self.ref_model = None + else: + ref_model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + ref_model_init_kwargs["device_map"] = None + ref_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + ref_model_path = get_config_model_id(self.model.config) + self.ref_model = create_model_from_path(ref_model_path, **ref_model_init_kwargs) + else: + self.ref_model = ref_model + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + if is_peft_model(self.model): + raise NotImplementedError( + "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " + "With PEFT, KTOTrainer does not keep a separate reference model in memory; instead, it recovers " + "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " + "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " + "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " + "feature request at https://github.com/huggingface/trl/issues." + ) + if args.precompute_ref_log_probs: + raise ValueError( + "You cannot use `sync_ref_model=True` together with `precompute_ref_log_probs=True`. " + "`precompute_ref_log_probs=True` assumes a fixed reference model, but with `sync_ref_model=True` " + "the reference model is periodically updated during training, making any precomputed reference " + "log-probs stale. Set `precompute_ref_log_probs=False` or disable `sync_ref_model`." + ) + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + self.use_liger_kernel = args.use_liger_kernel + # Import Liger kernel if enabled + if self.use_liger_kernel: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_kernel=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if self.loss_type in ["apo_zero_unpaired"]: + raise ValueError( + "You cannot set `loss_type='apo_zero_unpaired'` with liger-kernel." + "Only KTO loss is supported with liger-kernel." + ) + if self.precompute_ref_logps: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with liger kernel. Please set " + "`precompute_ref_log_probs=False`." + ) + if is_peft_model(self.model): + raise ValueError( + "You cannot use `use_liger_kernel=True` with Peft models. Please set `use_liger_kernel=False`." + ) + self.liger_loss_fn = LigerFusedLinearKTOLoss(beta=self.beta, use_ref_model=(self.ref_model is not None)) + + if self.precompute_ref_logps: + if isinstance(self.train_dataset, IterableDataset) or isinstance( + self.eval_dataset, (IterableDataset, IterableDatasetDict) + ): + raise ValueError( + "`precompute_ref_log_probs=True` is not supported with IterableDataset. Please use a map-style " + "Dataset or set `precompute_ref_log_probs=False`." + ) + self.train_dataset = self._precompute_ref_logps( + self.train_dataset, + "train", + self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size, + ) + if self.eval_dataset is not None: + if isinstance(self.eval_dataset, dict): + self.eval_dataset = { + name: self._precompute_ref_logps( + dataset, name, self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size + ) + for name, dataset in self.eval_dataset.items() + } + else: + self.eval_dataset = self._precompute_ref_logps( + self.eval_dataset, + "eval", + self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size, + ) + + def _tokenize( + self, + processing_class: PreTrainedTokenizerBase | ProcessorMixin, + input: str | list, + **kwargs, + ) -> dict[str, list]: + """Tokenize a single example for dataset preprocessing. + + Dispatches to `apply_chat_template` for conversational input (list of message dicts) and to `__call__` for + non-conversational input (str). + + Args: + processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`]): + The tokenizer or processor to use. + input (`str` or `list`): + A string for non-conversational input, or a list of message dicts for conversational input. + **kwargs: + Forwarded to `apply_chat_template` (e.g. `add_generation_prompt`, `return_assistant_tokens_mask`). + + Returns: + `dict` with at least an `"input_ids"` key mapping to a flat `list[int]`. + """ + if isinstance(input, list): # conversational: list of message dicts + if self._is_vlm: + input = prepare_multimodal_messages(input) + result = processing_class.apply_chat_template(input, tokenize=True, return_dict=True, **kwargs) + else: # non-conversational: plain text string + result = processing_class(text=input) + # VLMs emit a batch dimension even for single examples; unwrap it + if self._is_vlm: + return {k: v[0] for k, v in result.items()} + return result + + def _get_kl_dataset( + self, + dataset: Dataset | IterableDataset, + dataset_name: str, + args: KTOConfig, + ) -> Dataset | IterableDataset: + """ + Creates the KL dataset by creating mismatched (prompt, completion) pairs for KL divergence estimation. + + Args: + dataset (`Dataset` or `IterableDataset`): + Tokenized dataset with `prompt_ids` and `completion_ids` columns. + dataset_name (`str`): + Name used in progress bar descriptions. + args ([`KTOConfig`]): + Training arguments providing `per_device_train_batch_size` and `dataset_num_proc`. + + Returns: + `Dataset` or `IterableDataset` with a single `KL_completion_ids` column. + """ + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc or desc + map_kwargs["num_proc"] = args.dataset_num_proc + map_kwargs["desc"] = f"Extracting KL {dataset_name} dataset" + kl_dataset = dataset.map( + _get_kl_completion_ids, batched=True, batch_size=args.per_device_train_batch_size, **map_kwargs + ) + + def rename_kl_fn(example): + return {"KL_completion_ids": example["completion_ids"]} + + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Assembling KL {dataset_name} dataset" + column_names = get_dataset_column_names(dataset) + kl_dataset = kl_dataset.map( + rename_kl_fn, + remove_columns=[c for c in get_dataset_column_names(kl_dataset) if c in column_names], + **map_kwargs, + ) + return kl_dataset + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase | ProcessorMixin, + args: KTOConfig | None, + dataset_name: str, + ) -> Dataset | IterableDataset: + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed + first_example = next(iter(dataset)) + if "prompt" not in first_example: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Extracting prompt from {dataset_name} dataset" + dataset = dataset.map(extract_prompt, **map_kwargs) + + # Unpair the dataset if needed + first_example = next(iter(dataset)) + if "chosen" in first_example and "rejected" in first_example: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Unpairing {dataset_name} dataset" + dataset = unpair_preference_dataset(dataset, **map_kwargs) + + # Add EOS token if needed: non-conversational only + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if not example["completion"].endswith(eos_token): + example["completion"] = example["completion"] + eos_token + return example + + dataset = dataset.map(add_eos, fn_kwargs={"eos_token": self._tokenizer.eos_token}, **map_kwargs) + + # Tokenize dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + def tokenize_fn(example, processing_class): + if is_conversational(example): + chat_template_kwargs = example.get("chat_template_kwargs", {}) + prompt_ids = self._tokenize( + processing_class, + example["prompt"], + add_generation_prompt=True, + **chat_template_kwargs, + )["input_ids"] + prompt_completion_ids = self._tokenize( + processing_class, + example["prompt"] + example["completion"], + **chat_template_kwargs, + )["input_ids"] + else: + prompt_ids = self._tokenize(processing_class, example["prompt"])["input_ids"] + prompt_completion_ids = self._tokenize( + processing_class, example["prompt"] + example["completion"] + )["input_ids"] + + if not prompt_completion_ids[: len(prompt_ids)] == prompt_ids: + logger.warning( + "Mismatch between tokenized prompt and the start of tokenized prompt+completion. " + "This may be due to unexpected tokenizer behavior, whitespace issues, or special " + "token handling. Verify that the tokenizer is processing text consistently." + ) + + return { + "prompt_ids": prompt_ids, + "completion_ids": prompt_completion_ids[len(prompt_ids) :], + } + + dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) + + # Get KL datasets if needed + if self.calculate_KL: + # create pairs for estimating the KL term by flipping the matched pairs in each batch of size total_batch_size + # i.e., (x_1, y_1), ..., (x_n, y_n) --> (x_1, y_n), ..., (x_n, y_1) = (x'_1, y'_1), ..., (x'_n, y'_n) + kl_dataset = self._get_kl_dataset(dataset, dataset_name, args) + dataset = concatenate_datasets([dataset, kl_dataset], axis=1) + + # Calculate dataset desirability balance + if dataset_name == "train" and isinstance(dataset, Dataset): # IterableDataset does not support len + num_desirable = max(sum(dataset["label"]), 1) + num_undesirable = max(len(dataset["label"]) - num_desirable, 1) # "label" is binary + + if num_desirable != num_undesirable: + # The lower and upper bounds come from Eq. (8) of https://huggingface.co/papers/2402.01306 + des_weight_lower_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1, 2) + des_weight_upper_bound = round( + (num_undesirable * self.undesirable_weight / num_desirable) * 1.33, 2 + ) + und_weight_lower_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1.33, 2) + und_weight_upper_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1, 2) + + des_weight_in_range = des_weight_lower_bound <= self.desirable_weight <= des_weight_upper_bound + und_weight_in_range = und_weight_lower_bound <= self.undesirable_weight <= und_weight_upper_bound + + if not (des_weight_in_range or und_weight_in_range): + logger.warning( + "You have different amounts of desirable/positive and undesirable/negative examples but the " + "weights on the desirable and undesirable losses don't seem to be in an ideal range. Based " + f"on your data, we recommend EITHER " + f"desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}] or " + f"undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH). " + "See the documentation on how to optimally set these weights.", + ) + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). + if self._signature_columns is None: + if self._is_vision_dataset: + self._signature_columns = [ + "prompt", + "completion", + "image", + "images", + "label", + "chat_template_kwargs", + ] + else: + self._signature_columns = [ + "prompt_ids", + "completion_ids", + "KL_completion_ids", + "label", + "ref_logps", + "ref_KL_logps", + ] + + def _get_train_sampler(self, train_dataset: Dataset | None = None) -> Sampler | None: + if self.calculate_KL and Version(transformers.__version__) < Version("5.2.0"): + if train_dataset is None: + train_dataset = self.train_dataset + if train_dataset is None or not has_length(train_dataset): + return None + return SequentialSampler(train_dataset) + return super()._get_train_sampler( + train_dataset + ) # Override training step to add activation offloading context. + + def _precompute_ref_logps(self, dataset: Dataset, name: str, batch_size: int) -> Dataset: + model_hash = hash_module(self.ref_model or self.model) + fingerprint = Hasher.hash((dataset._fingerprint, model_hash, self.calculate_KL)) + cache_file = dataset._get_cache_file_path(fingerprint) + if os.path.exists(cache_file): + return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) + + dataloader = DataLoader( + dataset, + batch_size=batch_size, + collate_fn=self.data_collator, + num_workers=self.args.dataloader_num_workers, + pin_memory=self.args.dataloader_pin_memory, + shuffle=False, + ) + data_loader = self.accelerator.prepare(dataloader) + ref_logps = [] + ref_KL_logps = [] + for padded_batch in tqdm(iterable=data_loader, desc=f"Computing reference log probs for {name} dataset"): + ref_logp, ref_KL_logp = self.compute_ref_log_probs(padded_batch) + if self.calculate_KL: + ref_logp, ref_KL_logp = self.accelerator.gather_for_metrics((ref_logp, ref_KL_logp)) + ref_KL_logps.append(ref_KL_logp.cpu()) + else: + ref_logp = self.accelerator.gather_for_metrics(ref_logp) + ref_logps.append(ref_logp.cpu()) + + ref_logps = torch.cat(ref_logps) + if self.calculate_KL: + ref_KL_logps = torch.cat(ref_KL_logps) + + if self.accelerator.is_main_process: + + def add_ref_logps(batch, indices): + result = {"ref_logps": ref_logps[indices]} + if self.calculate_KL: + result.update({"ref_KL_logps": ref_KL_logps[indices]}) + return result + + dataset.map( + add_ref_logps, + with_indices=True, + batched=True, + remove_columns=dataset.column_names, + new_fingerprint=fingerprint, + desc=f"Caching reference log probs for {name} dataset", + ) + self.accelerator.wait_for_everyone() + + return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) + + def compute_ref_log_probs(self, inputs): + """Computes reference log probabilities for a single padded batch.""" + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + if self.ref_model is None: + if is_peft_model(self.model): + model = self.accelerator.unwrap_model(self.model) + with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): + completion_logits = self.model( + inputs["completion_input_ids"], + attention_mask=inputs["completion_attention_mask"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + inputs["KL_completion_input_ids"], + attention_mask=inputs["KL_completion_attention_mask"], + ).logits + else: + completion_logits = self.model( + inputs["completion_input_ids"], + attention_mask=inputs["completion_attention_mask"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + inputs["KL_completion_input_ids"], + attention_mask=inputs["KL_completion_attention_mask"], + ).logits + else: + completion_logits = self.ref_model( + inputs["completion_input_ids"], attention_mask=inputs["completion_attention_mask"] + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + inputs["KL_completion_input_ids"], + attention_mask=inputs["KL_completion_attention_mask"], + ).logits + + shift_logits = completion_logits[:, :-1, :] + per_token_logps = selective_log_softmax(shift_logits, inputs["completion_input_ids"][:, 1:]) + per_token_logps[inputs["completion_mask"][:, 1:] == 0] = 0.0 + completion_logps = per_token_logps.sum(-1) + + if self.calculate_KL: + shift_KL_logits = KL_logits[:, :-1, :] + KL_per_token_logps = selective_log_softmax(shift_KL_logits, inputs["KL_completion_input_ids"][:, 1:]) + KL_per_token_logps[inputs["KL_completion_mask"][:, 1:] == 0] = 0.0 + KL_logps = KL_per_token_logps.sum(-1) + else: + KL_logps = None + + return completion_logps, KL_logps + + def _compute_kl_logps(self, model, batch): + """Compute KL log probabilities for a given batch.""" + KL_logps = None + if self.calculate_KL: + _non_model_keys = { + "completion_input_ids", + "completion_attention_mask", + "completion_mask", + "KL_completion_mask", + "KL_completion_token_type_ids", + "KL_completion_mm_token_type_ids", + "label", + "ref_logps", + "ref_KL_logps", + } + KL_model_kwargs = {k: v for k, v in batch.items() if k not in _non_model_keys} + KL_model_kwargs["input_ids"] = KL_model_kwargs.pop("KL_completion_input_ids") + KL_model_kwargs["attention_mask"] = KL_model_kwargs.pop("KL_completion_attention_mask") + # KL sequences have different widths from the main completion after flush_left; override token-type + # tensors with the KL-specific ones the collator built for exactly this purpose. + if "KL_completion_token_type_ids" in batch: + KL_model_kwargs["token_type_ids"] = batch["KL_completion_token_type_ids"] + if "KL_completion_mm_token_type_ids" in batch: + KL_model_kwargs["mm_token_type_ids"] = batch["KL_completion_mm_token_type_ids"] + + with torch.no_grad(): + KL_logits = model(**KL_model_kwargs).logits + + shift_KL_logits = KL_logits[:, :-1, :] + KL_per_token_logps = selective_log_softmax(shift_KL_logits, batch["KL_completion_input_ids"][:, 1:]) + KL_per_token_logps[batch["KL_completion_mask"][:, 1:] == 0] = 0.0 + KL_logps = KL_per_token_logps.sum(-1) + return KL_logps + + def _compute_loss_liger(self, model, inputs, return_outputs): + if return_outputs: + raise RuntimeError( + "return_outputs=True is not supported with the Liger KTO loss. The Liger loss computes the loss " + "without materializing logits, so outputs cannot be returned." + ) + mode = "train" if self.model.training else "eval" + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in inputs.items()} + + labels = torch.tensor(batch["label"]) + num_chosen = labels.sum().to(self.accelerator.device) + num_rejected = (len(labels) - num_chosen).to(self.accelerator.device) + + policy_KL_logps = self._compute_kl_logps(model, batch) + ref_KL_logps = self._compute_kl_logps(self.ref_model, batch) + if self.calculate_KL: + kl = (policy_KL_logps - ref_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(self.accelerator.device) + + _non_model_keys = { + "completion_mask", + "KL_completion_input_ids", + "KL_completion_attention_mask", + "KL_completion_mask", + "KL_completion_token_type_ids", + "KL_completion_mm_token_type_ids", + "label", + "ref_logps", + "ref_KL_logps", + } + model_kwargs = {k: v for k, v in batch.items() if k not in _non_model_keys} + model_kwargs["input_ids"] = model_kwargs.pop("completion_input_ids") + model_kwargs["attention_mask"] = model_kwargs.pop("completion_attention_mask") + model_kwargs["use_cache"] = False + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + # `base_model` gives the inner module (skipping `lm_head`) — text decoder for LMs, multimodal wrapper for + # VLMs (so vision-token injection runs before the text decoder). `get_decoder()` won't do: on VLMs it + # returns just the text stack and feeds image-placeholder IDs through it. + # Pre-5.0 transformers VLMs set `base_model_prefix = ""` so `base_model is self` (re-runs `lm_head`). + # Fall back to `.model` there. + if self._is_vlm and Version(transformers.__version__) < Version("5.0.0"): + backbone, ref_backbone = model.model, self.ref_model.model + else: + backbone, ref_backbone = model.base_model, self.ref_model.base_model + + outputs = backbone(**model_kwargs) + + # reference model + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + ref_outputs = ref_backbone(**{k: v for k, v in model_kwargs.items() if k != "output_router_logits"}) + lm_head = model.get_output_embeddings() + ref_lm_head = self.ref_model.get_output_embeddings() + + shift_completion_mask = batch["completion_mask"][:, 1:] + target = batch["completion_input_ids"][:, 1:].clone() + target[shift_completion_mask == 0] = -100 + + ( + loss, + ( + chosen_logps_sum, + rejected_logps_sum, + chosen_logits_sum, + rejected_logits_sum, + chosen_rewards_sum, + rejected_rewards_sum, + ), + ) = self.liger_loss_fn( + _input=outputs.last_hidden_state[:, :-1], + lin_weight=lm_head.weight, + target=target, + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + preference_labels=torch.tensor(batch["label"], dtype=torch.bool).to(self.accelerator.device), + ref_input=ref_outputs.last_hidden_state[:, :-1], + ref_weight=ref_lm_head.weight, + ref_bias=ref_lm_head.bias if hasattr(lm_head, "bias") else None, + kl=kl, + ) + if self.aux_loss_enabled: + loss += self.aux_loss_coef * outputs.aux_loss + + self._metrics[mode]["kl"].append(kl.item()) + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + self._metrics[mode]["rewards/chosen"].append( + self.accelerator.gather_for_metrics(chosen_rewards_sum.nansum()).nansum().item() / all_num_chosen + ) + self._metrics[mode]["logps/chosen"].append( + self.accelerator.gather_for_metrics(chosen_logps_sum.nansum()).nansum().item() / all_num_chosen + ) + self._metrics[mode]["logits/chosen"].append( + self.accelerator.gather_for_metrics(chosen_logits_sum.nansum()).nansum().item() / all_num_chosen + ) + + if all_num_rejected > 0: + self._metrics[mode]["rewards/rejected"].append( + self.accelerator.gather_for_metrics(rejected_rewards_sum.nansum()).nansum().item() / all_num_rejected + ) + self._metrics[mode]["logps/rejected"].append( + self.accelerator.gather_for_metrics(rejected_logps_sum.nansum()).nansum().item() / all_num_rejected + ) + self._metrics[mode]["logits/rejected"].append( + self.accelerator.gather_for_metrics(rejected_logits_sum.nansum()).nansum().item() / all_num_rejected + ) + + if all_num_chosen > 0 and all_num_rejected > 0: + self._metrics[mode]["rewards/margins"].append( + self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1] + ) + + return loss + + def _compute_loss(self, model, inputs, return_outputs): + """Compute the KTO loss and other metrics for the given batch of inputs for train or test.""" + mode = "train" if self.model.training else "eval" + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in inputs.items()} + + labels = torch.tensor(batch["label"]) + num_chosen = labels.sum().to(self.accelerator.device) + num_rejected = (len(labels) - num_chosen).to(self.accelerator.device) + + policy_KL_logps = self._compute_kl_logps(model, batch) + + _non_model_keys = { + "completion_mask", + "KL_completion_input_ids", + "KL_completion_attention_mask", + "KL_completion_mask", + "KL_completion_token_type_ids", + "KL_completion_mm_token_type_ids", + "label", + "ref_logps", + "ref_KL_logps", + } + model_kwargs = {k: v for k, v in batch.items() if k not in _non_model_keys} + model_kwargs["input_ids"] = model_kwargs.pop("completion_input_ids") + model_kwargs["attention_mask"] = model_kwargs.pop("completion_attention_mask") + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model(**model_kwargs) + if self.aux_loss_enabled: + aux_loss = outputs.aux_loss + + shift_logits = outputs.logits[:, :-1, :] + per_token_logps = selective_log_softmax(shift_logits, batch["completion_input_ids"][:, 1:]) + per_token_logps[batch["completion_mask"][:, 1:] == 0] = 0.0 + completion_logps = per_token_logps.sum(-1) + + if completion_logps.shape[0] != len(batch["label"]): + raise ValueError( + "There is a mismatch between the number of examples in this batch and the number of " + "examples for which an output sequence was predicted." + ) + + device = outputs.logits.device + bool_labels = torch.as_tensor(batch["label"], dtype=torch.bool, device=device) + chosen_idx = torch.nonzero(bool_labels, as_tuple=False).view(-1) + rejected_idx = torch.nonzero(~bool_labels, as_tuple=False).view(-1) + + policy_chosen_logps = completion_logps.index_select(0, chosen_idx) + policy_rejected_logps = completion_logps.index_select(0, rejected_idx) + policy_chosen_logits = outputs.logits.index_select(0, chosen_idx) + policy_rejected_logits = outputs.logits.index_select(0, rejected_idx) + + if self.precompute_ref_logps: + ref_chosen_logps = batch["ref_logps"].index_select(0, chosen_idx) + ref_rejected_logps = batch["ref_logps"].index_select(0, rejected_idx) + if self.calculate_KL: + ref_KL_logps = batch["ref_KL_logps"] + else: + ref_KL_logps = None + else: + ref_model_kwargs = {k: v for k, v in model_kwargs.items() if k != "output_router_logits"} + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + if is_peft_model(self.model) and self.ref_model is None: + ref_model_unwrapped = self.accelerator.unwrap_model(self.model) + with use_adapter( + ref_model_unwrapped, adapter_name="ref" if "ref" in ref_model_unwrapped.peft_config else None + ): + ref_KL_logps = self._compute_kl_logps(self.model, batch) + ref_outputs = self.model(**ref_model_kwargs) + else: + ref_KL_logps = self._compute_kl_logps(self.ref_model, batch) + ref_outputs = self.ref_model(**ref_model_kwargs) + ref_shift_logits = ref_outputs.logits[:, :-1, :] + ref_per_token_logps = selective_log_softmax(ref_shift_logits, batch["completion_input_ids"][:, 1:]) + ref_per_token_logps[batch["completion_mask"][:, 1:] == 0] = 0.0 + ref_completion_logps = ref_per_token_logps.sum(-1) + ref_chosen_logps = ref_completion_logps.index_select(0, chosen_idx) + ref_rejected_logps = ref_completion_logps.index_select(0, rejected_idx) + + if self.calculate_KL: + kl = (policy_KL_logps - ref_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(policy_chosen_logps.device) + # Chosen losses + if policy_chosen_logps.shape[0] != 0 or ref_chosen_logps.shape[0] != 0: + chosen_logratios = policy_chosen_logps - ref_chosen_logps + + if self.loss_type == "kto": + # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306) + chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl)) + elif self.loss_type == "apo_zero_unpaired": + # Unpaired variant of Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are better than your model's default output + chosen_losses = 1 - F.sigmoid(self.beta * chosen_logratios) + + chosen_rewards = self.beta * chosen_logratios.detach() + + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + chosen_losses = torch.Tensor([]).to(self.accelerator.device) + chosen_rewards = torch.Tensor([]).to(self.accelerator.device) + # Rejected losses + if policy_rejected_logps.shape[0] != 0 or ref_rejected_logps.shape[0] != 0: + rejected_logratios = policy_rejected_logps - ref_rejected_logps + + if self.loss_type == "kto": + rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios)) + elif self.loss_type == "apo_zero_unpaired": + rejected_losses = F.sigmoid(self.beta * rejected_logratios) + + rejected_rewards = self.beta * rejected_logratios.detach() + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + rejected_losses = torch.Tensor([]).to(self.accelerator.device) + rejected_rewards = torch.Tensor([]).to(self.accelerator.device) + losses = torch.cat( + (self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses), + 0, + ) + + self._metrics[mode]["kl"].append(kl.item()) + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + self._metrics[mode]["rewards/chosen"].append( + self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() / all_num_chosen + ) + self._metrics[mode]["logps/chosen"].append( + self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() / all_num_chosen + ) + self._metrics[mode]["logits/chosen"].append( + self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() / all_num_chosen + ) + + if all_num_rejected > 0: + self._metrics[mode]["rewards/rejected"].append( + self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() / all_num_rejected + ) + self._metrics[mode]["logps/rejected"].append( + self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() / all_num_rejected + ) + self._metrics[mode]["logits/rejected"].append( + self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() / all_num_rejected + ) + + if all_num_chosen > 0 and all_num_rejected > 0: + self._metrics[mode]["rewards/margins"].append( + self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1] + ) + + loss = losses.nanmean() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return (loss, outputs) if return_outputs else loss + + def evaluate( + self, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> dict[str, float]: + # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way + # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer. `_prepare_dataset` is + # idempotent: it skips datasets that are already tokenized. A `str` selects a dataset that was already prepared + # at init time, so it's left untouched. + if not self._is_vision_dataset and eval_dataset is not None and not isinstance(eval_dataset, str): + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, self.processing_class, self.args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, self.processing_class, self.args, "eval") + # With `precompute_ref_log_probs`, `_compute_loss` reads the reference log-probs from the batch, so they + # must be precomputed here as well, mirroring `__init__`. + if self.precompute_ref_logps: + batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size + if isinstance(eval_dataset, dict): + eval_dataset = { + name: self._precompute_ref_logps(dataset, name, batch_size) + for name, dataset in eval_dataset.items() + } + else: + eval_dataset = self._precompute_ref_logps(eval_dataset, "eval", batch_size) + return super().evaluate( + eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix + ) + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + try: + if self.use_liger_kernel: + return self._compute_loss_liger(model, inputs, return_outputs) + return self._compute_loss(model, inputs, return_outputs) + except ValueError as e: + if "Image features and image tokens do not match" in str(e) and self.args.max_length is not None: + raise ValueError( + f"The current `max_length` ({self.args.max_length}) is too short and causes image placeholder " + f"tokens in `input_ids` to be truncated, while the corresponding image features remain intact. " + f"Please increase `max_length` or set it to `None` to disable truncation." + ) from e + raise + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and + # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(), self.compute_loss_context_manager(): + if prediction_loss_only: + loss = self.compute_loss(model, inputs, return_outputs=False) # logits aren't materialized with liger + logits, labels = None, None + else: + loss, outputs = self.compute_loss(model, inputs, return_outputs=True) + logits, labels = outputs.logits, inputs["completion_input_ids"] + return loss, logits, labels + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) +class UnslothKTOTrainer(_UnslothKTOTrainer): + """ + KTOTrainer(*args, **kwargs) + """ + def __init__( + self, + model, + ref_model = None, + args = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + data_collator = None, + callbacks = None, + peft_config = None, + compute_metrics = None, + **kwargs + ): + if args is None: args = UnslothKTOConfig() + use_bf16 = getattr(args, 'bf16', False) + if type(use_bf16) is not bool: use_bf16 = False + use_fp16 = getattr(args, 'fp16', False) + if type(use_fp16) is not bool: use_fp16 = False + force_float32 = False + full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' + if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): + print('Unsloth: Switching to float32 training since model cannot work with float16') + force_float32 = True + mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') + dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) + if dtype is None: dtype = model.get_input_embeddings().weight.dtype + from unsloth_zoo.utils import _get_dtype + dtype = _get_dtype(dtype) + float16 = dtype == torch.float16 + if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') + if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') + if force_float32: + # Forced float32 training + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': + # Mixed precision training + args.fp16 = float16 + args.bf16 = not float16 + os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' + # args.mixed_precision is a new argument which needs to be set now + elif mixed_precision_dtype == 'bfloat16': + # Both False since bfloat16 full finetuning doesn't do any autocasting. + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + + if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': + args.eval_strategy = 'steps' + if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 + ga_steps = getattr(args, 'gradient_accumulation_steps', None) + if ga_steps is not None and ga_steps > 1: + from transformers import __version__ as transformers_version + if Version(transformers_version) <= Version('4.45.2'): + print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' + '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') + if getattr(args, 'eval_strategy', 'no') != 'no': + eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) + if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size + if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps + fp16_full_eval = getattr(args, 'fp16_full_eval', False) + if type(fp16_full_eval) is not bool: fp16_full_eval = False + bf16_full_eval = getattr(args, 'bf16_full_eval', False) + if type(bf16_full_eval) is not bool: bf16_full_eval = False + if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True + if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False + if force_float32: + args.bf16_full_eval = False + args.fp16_full_eval = False + elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': + args.bf16_full_eval = True + args.fp16_full_eval = False + elif not bf16_full_eval and not fp16_full_eval: + args.bf16_full_eval = args.bf16 + args.fp16_full_eval = args.fp16 + _output_logits = False + if locals().get('compute_metrics', None) is not None: _output_logits = True + if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True + if _output_logits: + os.environ['UNSLOTH_RETURN_LOGITS'] = '1' + if model is not None: + _warnings_issued = getattr(model, 'warnings_issued', None) + if _warnings_issued is None: + model.warnings_issued = {} + elif not isinstance(_warnings_issued, dict): + try: + model.warnings_issued = dict(_warnings_issued) + except Exception: + model.warnings_issued = {} + if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): + pass + else: + model_max_seq_length = getattr(model, 'max_seq_length', None) + args_max_seq_length = getattr(args, 'max_seq_length', None) + if args_max_seq_length is None and model_max_seq_length is not None: + max_seq_length = model.max_seq_length + if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length + elif args_max_seq_length is not None and model_max_seq_length is not None: + if args_max_seq_length > model_max_seq_length: + print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' + 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') + args.max_seq_length = model_max_seq_length + if model is not None and hasattr(model, 'for_training'): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' + if 'processing_class' in locals(): + if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' + if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' + __tokenizer = processing_class if 'processing_class' in locals() else tokenizer + from unsloth_zoo.vision_utils import UnslothVisionDataCollator + if not isinstance(data_collator, UnslothVisionDataCollator): + if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names: + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names: + data_collator = DataCollatorForSeq2Seq( + __tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + else: + if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False + if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' + if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} + if not isinstance(data_collator, UnslothVisionDataCollator): + if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): + if isinstance(data_collator, DataCollatorForSeq2Seq): + data_collator = DataCollatorForSeq2Seq( + __tokenizer.tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer.tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('kto_trainer', other_metrics) + + # [TODO] Fix up DataParallel multiplying batch sizes + # [TODO] DDP works, but DP seems to not work? [TODO] + if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: + if getattr(args, "_n_gpu", 1) != 1: + args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + super().__init__( + model = model, + ref_model = ref_model, + args = args, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + data_collator = data_collator, + callbacks = callbacks, + peft_config = peft_config, + compute_metrics = compute_metrics,**kwargs) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() + if hasattr(self, 'neftune_hook_handle'): + self.neftune_hook_handle.remove() + if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle + if getattr(args, 'neftune_noise_alpha', None) is not None: + model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha + pass + if hasattr(self, 'accelerator'): + scaler = self.accelerator.scaler + current_model = model + while hasattr(current_model, 'model'): + current_model.accelerator_scaler = scaler + current_model = current_model.model + current_model.accelerator_scaler = scaler + pass + if hasattr(self, 'train'): + self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) + pass + if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): + _vllm_tok = self.llm.get_tokenizer() + _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) + if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: + _vllm_tok.chat_template = _pc.chat_template + pass + +pass + + +if hasattr(logger, "addFilter"): + import logging + class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not (self.text in x.getMessage()) + pass + logger.addFilter(HideLoggingMessage("`use_cache=True`")) + diff --git a/unsloth_compiled_cache/UnslothRLOOTrainer.py b/unsloth_compiled_cache/UnslothRLOOTrainer.py new file mode 100644 index 0000000000000000000000000000000000000000..f70a23f73741b1f2173fd87b512818c99b8b5af2 --- /dev/null +++ b/unsloth_compiled_cache/UnslothRLOOTrainer.py @@ -0,0 +1,2706 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.rloo_trainer import (Any, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, Dataset, DistributedBackend, GenerationConfig, IterableDataset, LoraConfig, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RLOOConfig, RLOOTrainer, RepeatSampler, RewardFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, _BaseTrainer, apply_chat_template, asyncio, atexit, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, disable_gradient_checkpointing, entropy_from_logits, gather, gather_object, get_config_model_id, get_peft_model, identity, inspect, is_conversational, is_peft_available, is_peft_model, is_rich_available, logger, math, nanmax, nanmin, nanstd, nn, np, pad, pd, peft, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, selective_log_softmax, set_seed, shuffle_sequence_dict, shutdown_event_loop_in_daemon, split_pixel_values_by_grid, split_tensor_dict, start_event_loop_in_daemon, textwrap, time, torch, transformers, unsplit_pixel_values_by_grid, unwrap_model_for_generation, use_adapter, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, Dataset, DistributedBackend, GenerationConfig, IterableDataset, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RLOOConfig, RLOOTrainer, RewardFunc, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, atexit, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, gather, get_config_model_id, get_peft_model, identity, inspect, is_peft_available, is_peft_model, logger, nn, np, pad, pd, peft, prepare_deepspeed, prepare_fsdp, set_seed, shutdown_event_loop_in_daemon, start_event_loop_in_daemon, time, torch, transformers, Any, np, profiling_decorator, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, PeftModel, PreTrainedModel, is_peft_available, logger, peft, torch) + + +import os +import math +import logging +from typing import * +from dataclasses import dataclass, field +from packaging.version import Version +import torch +import numpy as np +from contextlib import nullcontext +from torch.nn import functional as F +import inspect +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling +from transformers.training_args import ParallelMode +from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize + +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass +# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain +# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. +try: + from unsloth.models._utils import _unsloth_reset_stray_compile_cache +except Exception: + def _unsloth_reset_stray_compile_cache(self): pass +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Drop any torch.compile graph cache poisoned by a stray pre-train forward. + try: + _unsloth_reset_stray_compile_cache(self) + except Exception: + pass + # Finish the previous W&B run if this is a subsequent train() call. + # We do this at the START of train() (not the end) so that + # evaluate() / log() still work after train() completes. + # HF's WandbCallback.setup() will call wandb.init() for the new run. + # See: https://github.com/unslothai/unsloth/issues/3954 + if getattr(self, '_unsloth_training_completed', False): + try: + import wandb + if wandb.run is not None: + wandb.finish() + # Reset HF's WandbCallback so it calls wandb.init() for the new run + for cb in self.callback_handler.callbacks: + if type(cb).__name__ == 'WandbCallback': + cb._initialized = False + break + except: + pass + # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + output = f(self, *args, **kwargs) + # Restore previous mode when possible + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass + # Mark that training completed so the next train() call can + # finish this W&B run before starting a new one + self._unsloth_training_completed = True + return output + return wrapper +pass + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_hidden_states_selective_log_softmax( + hidden_states: torch.Tensor, + lm_head: torch.Tensor, + index: torch.Tensor, + chunks: int = 4, + logit_scale_multiply: float = 0.0, + logit_scale_divide: float = 0.0, + logit_softcapping: float = 0.0, + temperature: float = 1.0, +) -> torch.Tensor: + # All Unsloth Zoo code licensed under AGPL3 + flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + flat_index = index.reshape(-1) + + chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) + chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) + + all_per_token_logps = [] + + for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): + chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() + + if logit_scale_multiply != 0.0: + chunk_logits = chunk_logits * logit_scale_multiply + if logit_scale_divide != 0.0: + chunk_logits = chunk_logits / logit_scale_divide + if logit_softcapping != 0.0: + chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) + + chunk_logits = chunk_logits.to(torch.float32) + + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + + selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + + all_per_token_logps = torch.concat(all_per_token_logps) + + all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) + return all_per_token_logps + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_selective_log_softmax( + logits, + index, + temperature: float = 1.0, + chunks: int = 4, +): + chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) + chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) + all_per_token_logps = [] + # Per-chunk selective_log_softmax. + for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): + chunk_logits = chunk_logits.to(torch.float32) + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + pass + all_per_token_logps = torch.concat(all_per_token_logps) + all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) + return all_per_token_logps + +def calculate_pad_tokens_in_prompt( + input_ids: torch.Tensor, + logits_to_keep: int, + pad_token_id: int +) -> torch.Tensor: + """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" + if logits_to_keep >= input_ids.shape[1]: + raise ValueError("logits_to_keep must be smaller than the sequence length.") + + prompt_section = input_ids[:, :-logits_to_keep] + + padding_mask = (prompt_section == pad_token_id) + + pad_token_counts = padding_mask.sum(dim=1) + + return pad_token_counts + +def create_completion_attention_mask( + completion_input_ids: torch.Tensor, + left_pad_tokens_per_prompt: torch.Tensor, + max_left_pad: int, + pad_token_id: int +) -> torch.Tensor: + """Build a completion mask that zeros leading prompt and trailing pad tokens. + + For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) + this returns [0,0,0,1,1,1,0,0,0]. + """ + batch_size, completion_len = completion_input_ids.shape + device = completion_input_ids.device + + num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt + + indices = torch.arange(completion_len, device=device).unsqueeze(0) + shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) + + non_padding_mask = (completion_input_ids != pad_token_id) + + final_mask = shift_mask & non_padding_mask + + return final_mask + +def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: + """Move all padding tokens in each sequence to the right.""" + mask = (tensor != pad_id) + # stable=True since the binary mask is unordered. + sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) + packed_tensor = torch.gather(tensor, 1, sorted_indices) + return packed_tensor + +def align_logprobs_with_mask( + logprob_tensor: torch.Tensor, + attention_mask: torch.Tensor, + pad_value: float = 0.0 +) -> torch.Tensor: + """Align a log probability tensor with a given attention mask.""" + + device = logprob_tensor.device + batch_size, logprob_seq_len = logprob_tensor.shape + mask_seq_len = attention_mask.shape[1] + + padded_logprobs = torch.full( + attention_mask.shape, + fill_value=pad_value, + dtype=logprob_tensor.dtype, + device=device + ) + + left_pad_counts = torch.argmax(attention_mask, dim=1) + + cols = torch.arange(logprob_seq_len, device=device) + dest_indices = left_pad_counts.unsqueeze(1) + cols + + # Destination row indices, shape [batch_size, logprob_seq_len]. + row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) + + # Keep only in-bounds destinations, then scatter via advanced indexing. + valid_mask = dest_indices < mask_seq_len + valid_rows = row_indices[valid_mask] + valid_cols = dest_indices[valid_mask] + valid_vals = logprob_tensor[valid_mask] + padded_logprobs[valid_rows, valid_cols] = valid_vals + + return padded_logprobs + +def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" + if tool_mask is None: + return completion_mask + if tool_mask.shape[0] != completion_mask.shape[0]: + raise ValueError("tool_mask batch size must match completion_mask batch size.") + + tool_mask = tool_mask.to(device=completion_mask.device) + if tool_mask.shape == completion_mask.shape: + aligned_tool_mask = tool_mask + else: + aligned_tool_mask = align_logprobs_with_mask( + tool_mask, + completion_mask, + pad_value=0, + ) + return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) + +def autotune_batch_and_chunks( + total_input_rows, + seq_len, + hidden_size, + vocab_size, + dtype_bytes=16, + multiplier=None +): + if multiplier is None: + final_m = max(4, seq_len // 4096) + else: + final_m = multiplier + + if torch.cuda.is_available(): + free_bytes, _ = torch.cuda.mem_get_info() + limit_gb = (free_bytes / (1024**3))*.80 + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + # XPU: estimate free memory as total - reserved. + total_mem = torch.xpu.get_device_properties(0).total_memory + reserved_mem = torch.xpu.memory_reserved() + free_bytes = total_mem - reserved_mem + limit_gb = (free_bytes / (1024**3)) * 0.80 + else: + # Fallback: assume 8GB available. + limit_gb = 8.0 + + bytes_to_gb = 1024**3 + + b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) + + hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb + + base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb + logits_gb = base_logits / final_m + + total_mem_gb = hidden_gb + logits_gb + + valid_mask = total_mem_gb <= limit_gb + valid_indices = torch.nonzero(valid_mask, as_tuple=False) + + if valid_indices.shape[0] == 0: + #This means your GPU will OOM + return 4, final_m + + best_idx = valid_indices[0].item() + final_b = int(b_vals[best_idx].item()) + + return final_b, final_m + +def sanitize_logprob(logprob): + """Local port of trl.scripts.vllm_serve.sanitize_logprob. + Filters NaN logprobs from vLLM outputs.""" + value = logprob.logprob + if math.isnan(value): + logging.getLogger(__name__).warning( + f"Generated NaN logprob, token logprob '{logprob}' will be ignored" + ) + return None + return value +@dataclass +class UnslothRLOOConfig(RLOOConfig): + """ + + Configuration class for the [`RLOOTrainer`]. + + This class includes only the parameters that are specific to RLOO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`str`, `dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`RLOOTrainer`] is provided as a string. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to + [`~transformers.AutoModelForCausalLM.from_pretrained`] and + [`~transformers.AutoProcessor.from_pretrained`]. Also applied to reward-model and reward-tokenizer loads. + router_aux_loss_coef (`float`, *optional*, defaults to `0.001`): + Coefficient of the load-balancing auxiliary loss. Only has an effect when training a Mixture-of-Experts + (MoE) model; for other models it does nothing. The auxiliary loss is added to the training loss with this + weight. Set to `0.0` to disable it. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + + > Parameters that control the data preprocessing + + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + num_generations (`int`, *optional*, defaults to `2`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + num_generations_eval (`int` or `None`, *optional*): + Number of generations to sample during evaluation. This allows using fewer generations during evaluation to + save computation. If `None`, uses the value of `num_generations`. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + pad_to_multiple_of (`int`, *optional*): + If set, the prompts ids and completions ids will be padded to a multiple of this value. + + > Parameters that control generation + + generation_batch_size (`int`, *optional*): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one + generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. + steps_per_generation (`int`, *optional*): + Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive + with `generation_batch_size`. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int`, *optional*, defaults to `0`): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + generation_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or + `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the + generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict + with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. + chat_template_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + cache_implementation (`str`, *optional*): + Implementation of the cache method for faster generation when `use_vllm` is set to `False`. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"colocate"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + vllm_structured_outputs_regex (`str`, *optional*): + Regex for vLLM structured outputs. If `None` (default), structured outputs is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + vllm_group_port (`int`, *optional*, defaults to `51216`): + Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port + is occupied, there is no need to change it. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_max_model_length (`int`, *optional*): + Context window for vLLM. Set it to at least the maximum prompt length in the dataset plus + `max_completion_length`; if omitted, it is inferred from the model config. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but + waking the engine adds host–device transfer latency. + + > Parameters that control generation acceleration powered by transformers continuous batching + + use_transformers_continuous_batching (`bool`, *optional*, defaults to `False`): + Whether to use transformers' continuous batching engine for generating completions. Requires + `transformers>=5.8.0`. + transformers_continuous_batching_config (`dict`, *optional*): + Keyword arguments for [`~transformers.generation.ContinuousBatchingConfig`]. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.05`): + KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving training + speed. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + epsilon_high (`float`, *optional*): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + reward_weights (`list[float]`, *optional*): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + normalize_advantages (`bool`, *optional*, defaults to `False`): + Whether to normalize advantages. Normalization is done per generation batch to have mean `0.0` and standard + deviation of `1.0`. + reward_clip_range (`tuple[float, float]`, *optional*): + Clip range for rewards as (min, max). If `None`, no clipping is applied. + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or + `trackio`. + num_completions_to_print (`int`, *optional*): + Number of completions to print with `rich`. If `None`, all completions are logged. + log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts. If `True`, only unique prompts are logged. If `False`, all prompts are + logged. + + > Deprecated parameters + + use_transformers_paged: + + + + Parameter `use_transformers_paged` is deprecated and will be removed in version v2.0.0. Use + `use_transformers_continuous_batching` instead. + + + + > [!NOTE] + > These parameters have default values different from [`~transformers.TrainingArguments`]: + > - `logging_steps`: Defaults to `10` instead of `500`. + > - `gradient_checkpointing`: Defaults to `True` instead of `False`. + > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. + > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. + + """ + vllm_sampling_params: Optional[Any] = field( + default = None, + metadata = {'help': 'vLLM SamplingParams'}, + ) + unsloth_num_chunks : Optional[int] = field( + default = -1, + metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, + ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {'help': 'Multiplier for chunked logit computations.'}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, + ) + + def __init__( + self, + output_dir = None, + per_device_train_batch_size = 4, + num_train_epochs = 3.0, + max_steps = -1, + learning_rate = 5e-05, + lr_scheduler_type = 'linear', + lr_scheduler_kwargs = None, + warmup_steps = 0.1, + optim = 'adamw_8bit', + optim_args = None, + weight_decay = 0.001, + adam_beta1 = 0.9, + adam_beta2 = 0.999, + adam_epsilon = 1e-08, + optim_target_modules = None, + gradient_accumulation_steps = 2, + average_tokens_across_devices = True, + max_grad_norm = 1.0, + label_smoothing_factor = 0.0, + bf16 = False, + fp16 = False, + bf16_full_eval = False, + fp16_full_eval = False, + tf32 = None, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = None, + torch_compile = False, + torch_compile_backend = None, + torch_compile_mode = None, + use_liger_kernel = False, + liger_kernel_config = None, + use_cache = False, + neftune_noise_alpha = None, + torch_empty_cache_steps = 250, + auto_find_batch_size = False, + logging_strategy = 'steps', + logging_steps = 1, + logging_first_step = False, + log_on_each_node = True, + logging_nan_inf_filter = False, + include_num_input_tokens_seen = False, + log_level = 'passive', + log_level_replica = 'warning', + disable_tqdm = None, + report_to = 'none', + run_name = None, + project = 'huggingface', + trackio_space_id = 'trackio', + eval_strategy = 'no', + eval_steps = None, + eval_delay = 0, + per_device_eval_batch_size = 4, + prediction_loss_only = False, + eval_on_start = False, + eval_do_concat_batches = True, + eval_use_gather_object = False, + eval_accumulation_steps = 2, + batch_eval_metrics = False, + save_only_model = False, + save_strategy = 'steps', + save_steps = 500, + save_on_each_node = False, + save_total_limit = None, + enable_jit_checkpoint = False, + push_to_hub = False, + hub_token = None, + hub_private_repo = None, + hub_model_id = None, + hub_strategy = 'every_save', + hub_always_push = False, + hub_revision = None, + load_best_model_at_end = False, + metric_for_best_model = None, + greater_is_better = None, + ignore_data_skip = False, + restore_callback_states_from_checkpoint = False, + full_determinism = False, + seed = 3407, + data_seed = 3407, + use_cpu = False, + accelerator_config = None, + parallelism_config = None, + dataloader_drop_last = False, + dataloader_num_workers = 0, + dataloader_pin_memory = True, + dataloader_persistent_workers = False, + dataloader_prefetch_factor = None, + remove_unused_columns = False, + label_names = None, + train_sampling_strategy = 'random', + length_column_name = 'length', + ddp_find_unused_parameters = None, + ddp_bucket_cap_mb = None, + ddp_broadcast_buffers = None, + ddp_backend = None, + ddp_timeout = 1800, + fsdp = None, + fsdp_config = None, + deepspeed = None, + debug = '', + skip_memory_metrics = True, + do_train = False, + do_eval = False, + do_predict = False, + resume_from_checkpoint = None, + warmup_ratio = None, + logging_dir = None, + local_rank = -1, + model_init_kwargs = None, + trust_remote_code = False, + router_aux_loss_coef = 0.001, + disable_dropout = False, + num_generations = 8, + num_generations_eval = None, + max_completion_length = 256, + ds3_gather_for_generation = True, + shuffle_dataset = True, + pad_to_multiple_of = None, + generation_batch_size = None, + steps_per_generation = None, + temperature = 1.0, + top_p = 1.0, + top_k = None, + min_p = None, + generation_kwargs = {}, + chat_template_kwargs = None, + repetition_penalty = 1.0, + cache_implementation = None, + use_vllm = False, + vllm_mode = 'colocate', + vllm_model_impl = 'vllm', + vllm_enable_sleep_mode = False, + vllm_structured_outputs_regex = None, + vllm_server_base_url = None, + vllm_server_host = '0.0.0.0', + vllm_server_port = 8000, + vllm_server_timeout = 240.0, + vllm_group_port = 51216, + vllm_gpu_memory_utilization = 0.3, + vllm_max_model_length = None, + vllm_tensor_parallel_size = 1, + beta = 0.05, + num_iterations = 1, + epsilon = 0.2, + epsilon_high = None, + reward_weights = None, + normalize_advantages = False, + reward_clip_range = None, + mask_truncated_completions = False, + sync_ref_model = False, + ref_model_mixup_alpha = 0.6, + ref_model_sync_steps = 512, + log_completions = False, + num_completions_to_print = None, + log_unique_prompts = False, + use_transformers_continuous_batching = False, + transformers_continuous_batching_config = None, + use_transformers_paged = False, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, + + **kwargs, + ): + if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') + if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') + if num_train_epochs is None: + num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override + if output_dir is None and save_strategy == 'steps' and save_steps == 500: + output_dir = 'unsloth_training_checkpoints' + save_strategy = 'no' + if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': + from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION + if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: + from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE + pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE + + if steps_per_generation is None and generation_batch_size is None: + ga = gradient_accumulation_steps + world_size = int(os.environ.get('WORLD_SIZE', '1')) + if (ga * world_size * per_device_train_batch_size) % num_generations != 0: + print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) + per_device_train_batch_size = num_generations + + if temperature <= 0: + raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') + elif temperature >= 10: + raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') + + + super().__init__( + output_dir = output_dir, + per_device_train_batch_size = per_device_train_batch_size, + num_train_epochs = num_train_epochs, + max_steps = max_steps, + learning_rate = learning_rate, + lr_scheduler_type = lr_scheduler_type, + lr_scheduler_kwargs = lr_scheduler_kwargs, + warmup_steps = warmup_steps, + optim = optim, + optim_args = optim_args, + weight_decay = weight_decay, + adam_beta1 = adam_beta1, + adam_beta2 = adam_beta2, + adam_epsilon = adam_epsilon, + optim_target_modules = optim_target_modules, + gradient_accumulation_steps = gradient_accumulation_steps, + average_tokens_across_devices = average_tokens_across_devices, + max_grad_norm = max_grad_norm, + label_smoothing_factor = label_smoothing_factor, + bf16 = bf16, + fp16 = fp16, + bf16_full_eval = bf16_full_eval, + fp16_full_eval = fp16_full_eval, + tf32 = tf32, + gradient_checkpointing = gradient_checkpointing, + gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, + torch_compile = torch_compile, + torch_compile_backend = torch_compile_backend, + torch_compile_mode = torch_compile_mode, + use_liger_kernel = use_liger_kernel, + liger_kernel_config = liger_kernel_config, + use_cache = use_cache, + neftune_noise_alpha = neftune_noise_alpha, + torch_empty_cache_steps = torch_empty_cache_steps, + auto_find_batch_size = auto_find_batch_size, + logging_strategy = logging_strategy, + logging_steps = logging_steps, + logging_first_step = logging_first_step, + log_on_each_node = log_on_each_node, + logging_nan_inf_filter = logging_nan_inf_filter, + include_num_input_tokens_seen = include_num_input_tokens_seen, + log_level = log_level, + log_level_replica = log_level_replica, + disable_tqdm = disable_tqdm, + report_to = report_to, + run_name = run_name, + project = project, + trackio_space_id = trackio_space_id, + eval_strategy = eval_strategy, + eval_steps = eval_steps, + eval_delay = eval_delay, + per_device_eval_batch_size = per_device_eval_batch_size, + prediction_loss_only = prediction_loss_only, + eval_on_start = eval_on_start, + eval_do_concat_batches = eval_do_concat_batches, + eval_use_gather_object = eval_use_gather_object, + eval_accumulation_steps = eval_accumulation_steps, + batch_eval_metrics = batch_eval_metrics, + save_only_model = save_only_model, + save_strategy = save_strategy, + save_steps = save_steps, + save_on_each_node = save_on_each_node, + save_total_limit = save_total_limit, + enable_jit_checkpoint = enable_jit_checkpoint, + push_to_hub = push_to_hub, + hub_token = hub_token, + hub_private_repo = hub_private_repo, + hub_model_id = hub_model_id, + hub_strategy = hub_strategy, + hub_always_push = hub_always_push, + hub_revision = hub_revision, + load_best_model_at_end = load_best_model_at_end, + metric_for_best_model = metric_for_best_model, + greater_is_better = greater_is_better, + ignore_data_skip = ignore_data_skip, + restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, + full_determinism = full_determinism, + seed = seed, + data_seed = data_seed, + use_cpu = use_cpu, + accelerator_config = accelerator_config, + parallelism_config = parallelism_config, + dataloader_drop_last = dataloader_drop_last, + dataloader_num_workers = dataloader_num_workers, + dataloader_pin_memory = dataloader_pin_memory, + dataloader_persistent_workers = dataloader_persistent_workers, + dataloader_prefetch_factor = dataloader_prefetch_factor, + remove_unused_columns = remove_unused_columns, + label_names = label_names, + train_sampling_strategy = train_sampling_strategy, + length_column_name = length_column_name, + ddp_find_unused_parameters = ddp_find_unused_parameters, + ddp_bucket_cap_mb = ddp_bucket_cap_mb, + ddp_broadcast_buffers = ddp_broadcast_buffers, + ddp_backend = ddp_backend, + ddp_timeout = ddp_timeout, + fsdp = fsdp, + fsdp_config = fsdp_config, + deepspeed = deepspeed, + debug = debug, + skip_memory_metrics = skip_memory_metrics, + do_train = do_train, + do_eval = do_eval, + do_predict = do_predict, + resume_from_checkpoint = resume_from_checkpoint, + warmup_ratio = warmup_ratio, + logging_dir = logging_dir, + local_rank = local_rank, + model_init_kwargs = model_init_kwargs, + trust_remote_code = trust_remote_code, + router_aux_loss_coef = router_aux_loss_coef, + disable_dropout = disable_dropout, + num_generations = num_generations, + num_generations_eval = num_generations_eval, + max_completion_length = max_completion_length, + ds3_gather_for_generation = ds3_gather_for_generation, + shuffle_dataset = shuffle_dataset, + pad_to_multiple_of = pad_to_multiple_of, + generation_batch_size = generation_batch_size, + steps_per_generation = steps_per_generation, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + generation_kwargs = generation_kwargs, + chat_template_kwargs = chat_template_kwargs, + repetition_penalty = repetition_penalty, + cache_implementation = cache_implementation, + use_vllm = use_vllm, + vllm_mode = vllm_mode, + vllm_model_impl = vllm_model_impl, + vllm_enable_sleep_mode = vllm_enable_sleep_mode, + vllm_structured_outputs_regex = vllm_structured_outputs_regex, + vllm_server_base_url = vllm_server_base_url, + vllm_server_host = vllm_server_host, + vllm_server_port = vllm_server_port, + vllm_server_timeout = vllm_server_timeout, + vllm_group_port = vllm_group_port, + vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, + vllm_max_model_length = vllm_max_model_length, + vllm_tensor_parallel_size = vllm_tensor_parallel_size, + beta = beta, + num_iterations = num_iterations, + epsilon = epsilon, + epsilon_high = epsilon_high, + reward_weights = reward_weights, + normalize_advantages = normalize_advantages, + reward_clip_range = reward_clip_range, + mask_truncated_completions = mask_truncated_completions, + sync_ref_model = sync_ref_model, + ref_model_mixup_alpha = ref_model_mixup_alpha, + ref_model_sync_steps = ref_model_sync_steps, + log_completions = log_completions, + num_completions_to_print = num_completions_to_print, + log_unique_prompts = log_unique_prompts, + use_transformers_continuous_batching = use_transformers_continuous_batching, + transformers_continuous_batching_config = transformers_continuous_batching_config, + use_transformers_paged = use_transformers_paged,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier + + # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+ + if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: + if 'use_reentrant' in self.gradient_checkpointing_kwargs: + del self.gradient_checkpointing_kwargs['use_reentrant'] + +pass + +class _UnslothRLOOTrainer(_BaseTrainer): + """""" + + _tag_names = ["trl", "rloo"] + _name = "RLOO" + _paper = { + "title": "Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs", + "id": "2402.14740", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{ahmadian2024back, + title = {{Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs}}, + author = {Arash Ahmadian and Chris Cremer and Matthias Gall{\'{e}} and Marzieh Fadaee and Julia Kreutzer and Olivier Pietquin and Ahmet {\"{U}}st{\"{u}}n and Sara Hooker}, + year = 2024, + booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), {ACL} 2024, Bangkok, Thailand, August 11-16, 2024}, + pages = {12248--12267}, + publisher = {Association for Computational Linguistics}, + editor = {Lun{-}Wei Ku and Andre Martins and Vivek Srikumar}, + }"""), + } + + def __init__( + self, + model: "str | PreTrainedModel | PeftModel", + reward_funcs: RewardFunc | list[RewardFunc], + args: RLOOConfig | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + peft_config: "PeftConfig | None" = None, + ): + + if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): + if (getattr(args, 'use_vllm', False) == False): + args.use_vllm = True + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = RLOOConfig(f"{model_name}-RLOO") + + # Model + if isinstance(model, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + model = create_model_from_path(model, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `RLOOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do + _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + + # Some models [SmolVLM/Idefics3] don't support `logits_to_keep` argument and error out if we pass it + # Inspect the forward method before we wrap the model with PEFT + self.model_kwarg_keys = ( + inspect.signature(model.forward).parameters.keys() + if not hasattr(model, "get_base_model") + else inspect.signature(model.get_base_model().forward).parameters.keys() + ) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained( + get_config_model_id(model.config), + truncation_side="left", + padding_side="left", + trust_remote_code=args.trust_remote_code, + ) + + if args.use_transformers_continuous_batching and isinstance(processing_class, ProcessorMixin): + raise ValueError( + "`use_transformers_continuous_batching` does not support multimodal models. Use `use_vllm` instead." + ) + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + self._tokenizer = processing_class.tokenizer + elif isinstance(processing_class, PreTrainedTokenizerBase): + self._tokenizer = processing_class + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if self._tokenizer.pad_token is None: + self._tokenizer.pad_token = self._tokenizer.eos_token + + # PEFT + if False: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if is_peft_model(model): + raise ValueError( + "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " + "and unload the existing adapter, save the resulting base model, and then pass that base model along " + "with the new `peft_config` to the trainer." + ) + # Create PEFT model + # ZeRO-3 + PEFT for non-quantized models: + # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. + # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, + # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. + # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. + # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for + # non-quantized bf16 training. + # - See: + # - TRL issue: https://github.com/huggingface/trl/issues/6089 + # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 + # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg + get_peft_model_kwargs = {} + if ( + args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and not _is_quantized_model + and Version(peft.__version__) >= Version("0.12.0") + ): + get_peft_model_kwargs["autocast_adapter_dtype"] = False + model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + + elif is_peft_model(model): + # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy + # of the "default" adapter, so that we can use it as the reference model during the training. PEFT only + # supports one adapter per model when the LoRA config uses `target_parameters` [see peft#3340], so in that + # case we skip the "ref" adapter and compute the reference log probs with adapters disabled, i.e. with the + # base model. + default_config = model.peft_config["default"] + if isinstance(default_config, LoraConfig) and default_config.target_parameters: + logger.warning( + "PEFT can't add a frozen reference adapter alongside one that uses `target_parameters` " + "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " + "If you wrapped the model only to apply LoRA, pass a `peft_config` to the trainer instead; if you " + "wrapped it deliberately (pretrained adapter or custom init), note that the base model matches " + "your adapter only when it's freshly zero-initialized. If it is, this warning is safe to ignore." + ) + else: + model.add_adapter("ref", default_config) + for name, param in model.named_parameters(): + if ".default." in name: + ref_name = name.replace(".default.", ".ref.") + ref_param = model.get_parameter(ref_name) + ref_param.data.copy_(param.data) + + # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally + # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 + if is_peft_model(model) and args.gradient_checkpointing: + model.enable_input_require_grads() + + # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the + # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by + # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for + # quantized models. See: https://github.com/huggingface/peft/issues/2889 + if _is_quantized_model: + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.bfloat16) + + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + self._has_async_funcs = any(inspect.iscoroutinefunction(func) for func in self.reward_funcs) + if self._has_async_funcs: + self.async_loop_thread, self.async_loop, self.async_loop_ready_event = start_event_loop_in_daemon( + name="RLOOTrainer-AsyncRewardLoop" + ) + # wait until the event loop is running in the daemon thread + self.async_loop_ready_event.wait() + atexit.register(shutdown_event_loop_in_daemon, self.async_loop_thread, self.async_loop) + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError( + f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " + f"reward functions ({len(reward_funcs)})." + ) + + for i, (reward_processing_class, reward_func) in enumerate( + zip(reward_processing_classes, reward_funcs, strict=True) + ): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained( + get_config_model_id(reward_func.config), trust_remote_code=args.trust_remote_code + ) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + + self.reward_processing_classes = reward_processing_classes + + # Training arguments + self.max_completion_length = args.max_completion_length + self.num_generations = args.num_generations + self.num_generations_eval = args.num_generations_eval or self.num_generations + self.chat_template_kwargs = args.chat_template_kwargs or {} + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_transformers_continuous_batching = args.use_transformers_continuous_batching + if self.use_transformers_continuous_batching: + if not Version(transformers.__version__) >= Version("5.8.0"): + raise ImportError( + "Using `use_transformers_continuous_batching` requires transformers>=5.8.0. " + "Please upgrade with `pip install --upgrade transformers`." + ) + from transformers.generation import ContinuousBatchingConfig + + cb_kwargs = dict(args.transformers_continuous_batching_config or {}) + # The transformers default [0.9] leaves almost no VRAM for the training backward pass; + # use a training-aware default unless the user has set it explicitly. + cb_kwargs.setdefault("max_memory_percent", 0.5) + self.continuous_batching_config = ContinuousBatchingConfig(**cb_kwargs) + else: + self.continuous_batching_config = None + self.pad_to_multiple_of = args.pad_to_multiple_of + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.normalize_advantages = args.normalize_advantages + self.mask_truncated_completions = args.mask_truncated_completions + self.reward_clip_range = args.reward_clip_range + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + + if train_dataset is None: + raise ValueError("`train_dataset` is required") + elif ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in RLOOTrainer. Please use a standard dataset instead." + ) + + # Multi-step + self.num_iterations = args.num_iterations + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + + # MoE load-balancing auxiliary loss, applied to Mixture-of-Experts models [no effect otherwise] + text_config = model.config.get_text_config() + is_moe = getattr(text_config, "output_router_logits", None) is not None + self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0 + self.router_aux_loss_coef = args.router_aux_loss_coef + # Tracks the number of iterations [forward + backward passes], including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in RLOO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + ) + + # Reference model + self.beta = args.beta + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + self.ref_model = create_model_from_path(get_config_model_id(self.model.config), **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self._current_train_step_time = 0.0 + self.log_completions = args.log_completions + self.log_unique_prompts = args.log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # Keep logs sized to the generation batch to record only outputs from the latest model update. + self._logs = { + "images": deque(maxlen=args.generation_batch_size), + "prompt": deque(maxlen=args.generation_batch_size), + "completion": deque(maxlen=args.generation_batch_size), + "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + "advantages": deque(maxlen=args.generation_batch_size), + "extra": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + } + # Buffers for user-logged data from reward functions, flushed after gathering + self._pending_extra_logs = defaultdict(list) + self._pending_metrics = defaultdict(list) + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + self.vllm_generation = VLLMGeneration( + model=self.model, + accelerator=self.accelerator, + processing_class=self.processing_class, + mode=args.vllm_mode, + structured_outputs_regex=args.vllm_structured_outputs_regex, + server_base_url=args.vllm_server_base_url, + server_host=args.vllm_server_host, + server_port=args.vllm_server_port, + group_port=args.vllm_group_port, + server_timeout=args.vllm_server_timeout, + tensor_parallel_size=args.vllm_tensor_parallel_size, + gpu_memory_utilization=args.vllm_gpu_memory_utilization, + max_model_length=args.vllm_max_model_length, + max_num_seqs=args.per_device_train_batch_size + * args.vllm_tensor_parallel_size + * args.steps_per_generation, + enable_sleep_mode=args.vllm_enable_sleep_mode, + model_impl=args.vllm_model_impl, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + min_p=self.min_p, + max_completion_length=self.max_completion_length, + logprobs=None, + generation_kwargs=args.generation_kwargs, + ) + self._last_loaded_step = -1 + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": self._tokenizer.pad_token_id, + "bos_token_id": self._tokenizer.bos_token_id, + "eos_token_id": self._tokenizer.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs, disable_compile=True) + # Keep training-specific generation kwargs to overwrite model's original generation config + self.generation_kwargs = generation_kwargs + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + self._dist = DistributedBackend(self.accelerator) + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + if self.beta == 0.0: + raise ValueError( + "You passed `sync_ref_model=True` while `beta=0.0`, which means the reference model is not used " + "during training. Consequently, RLOOTrainer does not create a `ref_model` instance, and there is " + "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." + ) + if is_peft_model(model): + raise NotImplementedError( + "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " + "With PEFT, RLOOTrainer does not keep a separate reference model in memory; instead, it recovers " + "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " + "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " + "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " + "feature request at https://github.com/huggingface/trl/issues." + ) + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). In RLOOTrainer, we preprocess data, so using the model's signature columns doesn't + # work. Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt", "image", "images"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. + def get_train_dataloader(self): + return self._get_dataloader( + dataset=self.train_dataset, + description="Training", + batch_size=self._train_batch_size * self.args.steps_per_generation, # < this is the change + sampler_fn=self._get_train_sampler, + is_training=True, + ) + + def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. Each row shows the per-step batch + # returned by `_prepare_inputs`; rows within a `steps_per_generation` block are slices of the same + # generated batch. When `num_iterations > 1`, that block is reused for multiple optimization passes + # before regenerating. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations_eval, + seed=self.args.seed, + ) + + @profiling_decorator + def _get_per_token_logps_and_entropies( + self, + model, + input_ids, + attention_mask, + logits_to_keep, + batch_size=None, + compute_entropy=False, + compute_aux_loss=False, + pixel_values=None, + image_grid_thw=None, + num_images=None, + pixel_attention_mask=None, + spatial_shapes=None, + num_tiles=None, + image_sizes=None, + token_type_ids=None, + mm_token_type_ids=None, + image_position_ids=None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Compute log-probs, (optionally) entropies, and (optionally) the MoE load-balancing aux loss.""" + batch_size = batch_size or input_ids.size(0) # Chunk inputs into smaller batches to reduce memory peak + all_logps = [] + all_entropies = [] + all_aux_losses = [] + for start in range(0, input_ids.size(0), batch_size): + input_ids_batch = input_ids[start : start + batch_size] + attention_mask_batch = attention_mask[start : start + batch_size] + + # Build model inputs + model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} + if image_grid_thw is not None and pixel_values is not None: + rows_per_image = image_grid_thw.prod(dim=-1) + rows_per_sample = torch.split(rows_per_image, num_images) + rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) + cum_rows = torch.cat([torch.tensor([0], device=rows_per_sample.device), rows_per_sample.cumsum(0)]) + row_start, row_end = cum_rows[start].item(), cum_rows[start + batch_size].item() + model_inputs["pixel_values"] = pixel_values[row_start:row_end] + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] + model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end] + elif image_position_ids is not None and pixel_values is not None: + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] + model_inputs["pixel_values"] = pixel_values[img_start:img_end] + model_inputs["image_position_ids"] = image_position_ids[img_start:img_end] + elif spatial_shapes is not None and pixel_values is not None: + # LFM2-VL tensors are tile-indexed. + cum_tiles = torch.tensor([0] + num_tiles).cumsum(0) + tile_start, tile_end = cum_tiles[start], cum_tiles[start + batch_size] + model_inputs["pixel_values"] = pixel_values[tile_start:tile_end] + model_inputs["pixel_attention_mask"] = pixel_attention_mask[tile_start:tile_end] + model_inputs["spatial_shapes"] = spatial_shapes[tile_start:tile_end] + elif pixel_values is not None: + model_inputs["pixel_values"] = pixel_values[start : start + batch_size] + if pixel_attention_mask is not None and spatial_shapes is None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask[start : start + batch_size] + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes[start : start + batch_size] + if token_type_ids is not None: + model_inputs["token_type_ids"] = token_type_ids[start : start + batch_size] + if mm_token_type_ids is not None: + model_inputs["mm_token_type_ids"] = mm_token_type_ids[start : start + batch_size] + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + # MoE models: request router logits so the model returns `outputs.aux_loss`. VLM wrappers honor this only + # as a forward kwarg (not from the model config), so it must be passed here. + if compute_aux_loss: + model_inputs["output_router_logits"] = True + + outputs = model(**model_inputs) + logits = outputs.logits + # Exclude the last value: it corresponds to the next token pred + logits = logits[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + logits = logits[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + # Divide logits by sampling temperature. + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + logits.div_(self.temperature) + completion_ids = input_ids_batch[:, -logits_to_keep:] + logps = selective_log_softmax(logits, completion_ids) # compute logprobs + all_logps.append(logps) + + if compute_entropy: + with torch.no_grad(): + entropies = entropy_from_logits(logits) + all_entropies.append(entropies) + + if compute_aux_loss: + all_aux_losses.append(outputs.aux_loss) + + logps = torch.cat(all_logps, dim=0) + entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None + aux_loss = torch.stack(all_aux_losses).mean() if compute_aux_loss else None + return logps, entropies, aux_loss + + def training_step(self, model, inputs, num_items_in_batch): + time_before = time.perf_counter() + output = super().training_step(model, inputs, num_items_in_batch) + self._step += 1 + time_after = time.perf_counter() + self._current_train_step_time += time_after - time_before + if self._step % self.current_gradient_accumulation_steps == 0: + self._metrics["train"]["step_time"].append(self._current_train_step_time) + self._current_train_step_time = 0.0 + return output + + @profiling_decorator + def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = split_pixel_values_by_grid(generation_batch) + + try: generation_batch = shuffle_sequence_dict(generation_batch) + + except: pass + generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) + self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + return inputs + + def _log_completion_extra(self, column: str, values: list): + """ + Log extra columns to the completions table. Called from reward functions via the `log_extra` kwarg. + + Args: + column (`str`): + Name of the column to add. + values (`list`): + Values for the column, one per sample in the batch. + """ + self._pending_extra_logs[column].extend(values) + + def _log_metric(self, name: str, value: float): + """ + Log a scalar metric from a reward function. Called via the `log_metric` kwarg. Values are averaged over each + logging step and reported alongside built-in metrics like `kl` and `entropy`. + + Args: + name (`str`): + Name of the metric. + value (`float`): + Scalar value for this batch. + """ + self._pending_metrics[name].append(value) + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + # This allows for dynamic reward shaping based on training progress. + reward_kwargs["trainer_state"] = self.state + + # Allow reward functions to log extra columns to the completions table. + reward_kwargs["log_extra"] = self._log_completion_extra + + # Allow reward functions to log additional scalar metrics. + reward_kwargs["log_metric"] = self._log_metric + + async_funcs_info = [] # async custom functions for asyncio.gather + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) + ): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + with profiling_context(self, reward_func_name): + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] + texts = [ + apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] + for x in messages + ] + else: + texts = [p + c for p, c in zip(prompts, completions, strict=True)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + elif inspect.iscoroutinefunction(reward_func): # Separate async reward funcs to run them in parallel later + async_funcs_info.append((i, reward_func, reward_func_name)) + else: + # Run synchronous reward function + with profiling_context(self, reward_func_name): + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # Execute async custom functions in parallel using asyncio.gather + if async_funcs_info: + + async def _invoke_async(index, func, func_name): + with profiling_context(self, func_name): + output = await func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + output = [r if r is not None else torch.nan for r in output] + return index, output + + async def _run_async_funcs(): + coros = [_invoke_async(i, func, func_name) for (i, func, func_name) in async_funcs_info] + return await asyncio.gather(*coros) + + async_results = asyncio.run_coroutine_threadsafe(_run_async_funcs(), self.async_loop).result() + for idx, output_reward_func in async_results: + rewards_per_func[:, idx] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = { + key: value[nan_row_idx] + for key, value in reward_kwargs.items() + if key not in ("trainer_state", "log_extra", "log_metric") + } + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + logger.warning( + f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _tokenize_prompts(self, prompts: list): + """Tokenize prompts and extract images/multimodal fields for generation.""" + if is_conversational({"prompt": prompts[0]}): + # Extract images from messages for VLM support + images = [] + has_images = False + for prompt in prompts: + prompt_images = [] + for message in prompt: + if isinstance(message["content"], list): + for part in message["content"]: + if part["type"] == "image": + prompt_images.append(part["image"]) + has_images = True + images.append(prompt_images if prompt_images else None) + images = images if has_images else None + + # Workaround for a bug in transformers 5.3.0 where some processors (e.g. Qwen2.5-VL) crash on + # batched unpadded input (transformers#44514). + # Fixed in transformers 5.4.0 (transformers#44563). + needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") + tokenized = self.processing_class.apply_chat_template( + conversation=prompts, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **({"padding": True} if needs_padding_workaround else {}), + **self.chat_template_kwargs, + ) + if needs_padding_workaround: + # Unpad input_ids: remove padding tokens using attention_mask to get per-sequence lists + prompt_ids = [ + [tok for tok, m in zip(ids, mask, strict=True) if m] + for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) + ] + else: + prompt_ids = tokenized["input_ids"] + # For VLMs, the processor returns extra multimodal fields (pixel_values, image_grid_thw, etc.) + multimodal_fields = {k: v for k, v in tokenized.items() if k not in ("input_ids", "attention_mask")} + else: + prompt_ids = self.processing_class(text=prompts)["input_ids"] + images = None + multimodal_fields = {} + return prompt_ids, images, multimodal_fields + + def _generate_single_turn(self, prompt_ids, images, multimodal_fields): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + # Sync weights if training step changed + if self.state.global_step != self._last_loaded_step: + with profiling_context(self, "sync_weights"): + self.vllm_generation.sync_weights() + self._last_loaded_step = self.state.global_step + + # Generate using vLLM (note: RLOO doesn't use logprobs from generation, so we ignore them) + num_generations = self.num_generations if mode == "train" else self.num_generations_eval + _, completion_ids, _, _ = self.vllm_generation.generate( + prompts=prompt_ids, + images=images, + num_generations=num_generations, + profiler=profiling_context(self, "vLLM.generate"), + ) + + elif self.use_transformers_continuous_batching: + with ( + profiling_context(self, "transformers.generate_batch"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + self._dist.summon_full_params(self.model_wrapped, recurse=False), + ): + # Cast to the appropriate dtype based on training configuration + if self.args.bf16: + unwrapped_model.to(torch.bfloat16) + elif self.args.fp16: + unwrapped_model.to(torch.float16) + all_outputs = unwrapped_model.generate_batch( + prompt_ids, + generation_config=self.generation_config, + continuous_batching_config=self.continuous_batching_config, + progress_bar=False, + ) + unwrapped_model.train() + completion_ids = [output.generated_tokens for output in all_outputs.values()] + + else: + # Regular generation path: left-pad token IDs into tensors + prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] + padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left") + attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left") + generate_inputs = {"input_ids": padded_ids, "attention_mask": attention_mask} + # For VLMs, include multimodal fields as tensors (pixel_values, image_grid_thw, etc.) + for k, v in multimodal_fields.items(): + if isinstance(v, torch.Tensor): + generate_inputs[k] = v + elif isinstance(v, list) and v and isinstance(v[0], list): + # Per-token field (e.g., token_type_ids): left-pad like input_ids + generate_inputs[k] = pad([torch.tensor(x) for x in v], padding_value=0, padding_side="left") + else: + generate_inputs[k] = torch.tensor(np.array(v)) + generate_inputs = super()._prepare_inputs(generate_inputs) + + with ( + profiling_context(self, "transformers.generate"), + unwrap_model_for_generation( + self.model_wrapped, + self.accelerator, + gather_deepspeed3_params=self.args.ds3_gather_for_generation, + generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 + ) as unwrapped_model, + torch.no_grad(), + self._dist.summon_full_params(self.model_wrapped, recurse=False), + ): + prompt_completion_ids = unwrapped_model.generate( + **generate_inputs, generation_config=self.generation_config + ) + # Compute prompt length and extract completion ids + prompt_length = generate_inputs["input_ids"].size(1) + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self._tokenizer.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + completion_ids = [ + c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) + ] + + return completion_ids + + def _generate(self, prompts: list): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + # Copy the prompts to avoid modifying the original list + prompts = copy.deepcopy(prompts) + + prompt_ids, images, multimodal_fields = self._tokenize_prompts(prompts) + completion_ids = self._generate_single_turn(prompt_ids, images, multimodal_fields) + + # Decode completions. It's important to use `parse_response` when possible, because it handles tool calls. + if is_conversational({"prompt": prompts[0]}): + contents = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + completions = [[{"role": "assistant", "content": content}] for content in contents] + else: + completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + + # Get completion length per sequence, used for logging + prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) + completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) + agg_prompt_lengths = self.accelerator.gather(prompt_lengths) + agg_completion_lengths = self.accelerator.gather(completion_lengths) + total_prompt_tokens = agg_prompt_lengths.sum() + total_completion_tokens = agg_completion_lengths.sum() # = num_items_in_batch, required for the DAPO loss + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += (total_prompt_tokens + total_completion_tokens).item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + eos_and_pad = [self._tokenizer.eos_token_id, self._tokenizer.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) + agg_is_truncated = self.accelerator.gather(is_truncated) + self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) + term_completion_lengths = agg_completion_lengths[~agg_is_truncated] + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + return prompt_ids, completion_ids, completions + + def _generate_and_score_completions( + self, inputs: list[dict[str, torch.Tensor | Any]] + ) -> dict[str, torch.Tensor | Any]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + + if "images" in inputs[0]: + images = [example.get("images") for example in inputs] + elif "image" in inputs[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] + else: + images = None + # Transformers requires at least one image in the batch, otherwise it throws an error + if images is not None and all(img_list == [] for img_list in images): + images = None + + # If the prompts are conversational and the inputs contain images, we need to convert the prompts from + # [{"role": "user", "content": "What color is the sky?"}] to + # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] + if images is not None: + if not is_conversational(inputs[0]): + raise ValueError( + "Multimodal training requires conversational prompts. It looks like the dataset contains " + "non-conversational inputs, likely because a chat template was applied before passing the dataset " + "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat " + "template internally." + ) + prompts = [ + prepare_multimodal_messages(prompt, images=image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ] + + prompt_ids_list, completion_ids_list, completions = self._generate(prompts) + + # Convert lists of token IDs to padded tensors + prompt_ids = [torch.tensor(ids) for ids in prompt_ids_list] + prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] + prompt_ids = pad( + prompt_ids, + padding_value=self._tokenizer.pad_token_id, + padding_side="left", + pad_to_multiple_of=self.pad_to_multiple_of, + ).to(device=device) + prompt_mask = pad( + prompt_mask, padding_value=0, padding_side="left", pad_to_multiple_of=self.pad_to_multiple_of + ).to(device=device) + completion_ids = [torch.tensor(ids) for ids in completion_ids_list] + completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] + completion_ids = pad( + completion_ids, + padding_value=self._tokenizer.pad_token_id, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ).to(device=device) + completion_mask = pad( + completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ).to(device=device) + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + eos_and_pad = [self._tokenizer.eos_token_id, self._tokenizer.pad_token_id] + # Mask completion_mask for attention masking + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) + completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + num_images = [len(img_list) if img_list else 0 for img_list in images] if images is not None else None + + # Get forward_kwargs for models with multimodal inputs + if images is not None: + prompts_text = [ + apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] + for prompt in prompts + ] + prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") + prompt_inputs = super()._prepare_inputs(prompt_inputs) + forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} + else: + forward_kwargs = {} + + # Recover LFM2-VL tile counts; the full processor drops row/column metadata. + num_tiles = None + if images is not None and "spatial_shapes" in forward_kwargs: + image_info = self.processing_class.image_processor( + images=images, return_tensors="pt", return_row_col_info=True + ) + tiles_per_image = image_info["image_rows"] * image_info["image_cols"] + if self.processing_class.image_processor.use_thumbnail: + tiles_per_image = tiles_per_image + (tiles_per_image > 1).to(tiles_per_image.dtype) + num_tiles = [group.sum().item() for group in torch.split(tiles_per_image, num_images)] + + # If token_type_ids are used, extend them with zeros for the completion part + if "token_type_ids" in forward_kwargs: + token_type_ids = forward_kwargs["token_type_ids"] + if self.pad_to_multiple_of is not None: + # Needed only with pad_to_multiple_of: otherwise prompt_ids and token_type_ids must have equal len + padding_size = prompt_ids.size(1) - token_type_ids.size(1) + if padding_size > 0: + token_type_ids = torch.cat( + [token_type_ids.new_zeros((token_type_ids.size(0), padding_size)), token_type_ids], dim=1 + ) + forward_kwargs["token_type_ids"] = torch.cat( + [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + # If mm_token_type_ids are used, extend them with zeros for the completion part + if "mm_token_type_ids" in forward_kwargs: + mm_token_type_ids = forward_kwargs["mm_token_type_ids"] + if self.pad_to_multiple_of is not None: + # Needed only with pad_to_multiple_of: otherwise prompt_ids and mm_token_type_ids must have equal len + padding_size = prompt_ids.size(1) - mm_token_type_ids.size(1) + if padding_size > 0: + mm_token_type_ids = torch.cat( + [mm_token_type_ids.new_zeros((mm_token_type_ids.size(0), padding_size)), mm_token_type_ids], + dim=1, + ) + forward_kwargs["mm_token_type_ids"] = torch.cat( + [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + + # When gradient checkpointing is enabled with use_reentrant=True (non default), calling the model inside a + # torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True"). + # Temporarily disable checkpointing to avoid this warning during inference. + with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): + # Compute the per-token log probabilities for the current model + old_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size, + num_images=num_images, + num_tiles=num_tiles, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids + ) + old_logps = (old_per_token_logps * completion_mask).sum(1) # mask out padding and tokens after EOS + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( + self.ref_model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + num_tiles=num_tiles, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids + ) + else: + # When training a PEFT adapter, how we obtain the reference depends on the setup: + # - New adapter: disabling adapters yields the base model. + # - Re-training an existing adapter: an initial copy is loaded under the name "ref". + model = self.accelerator.unwrap_model(self.model) + with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): + ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + num_tiles=num_tiles, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids + ) + else: + ref_per_token_logps = None + + # Decode + prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + num_generations = self.num_generations if mode == "train" else self.num_generations_eval + + # A completion for which every reward function returned None is unscorable. nansum would collapse it to 0, + # which both biases the leave-one-out baseline and hands the completion a spurious advantage. Mark these rows + # NaN so they're excluded from the (nan-aware) baseline below; their advantage is forced to 0 afterwards. + unscorable_mask = torch.isnan(rewards_per_func).all(dim=1) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + rewards[unscorable_mask] = torch.nan + + # Apply reward clipping if specified + if self.reward_clip_range: + rewards = rewards.clamp(min=self.reward_clip_range[0], max=self.reward_clip_range[1]) + + # Include the KL penalty in the reward + if self.beta != 0.0: + # RLOO uses the first-order log ratio for the per-token KL estimate, following the original RLOO paper + # (Ahmadian et al., 2024, https://huggingface.co/papers/2405.14782). Unlike GRPOTrainer's Schulman + # approximation (always >= 0), this can be negative per token. The divergence is intentional: RLOO applies + # KL as a reward penalty (summed across tokens per sequence), while GRPO adds it to the per-token loss. + per_token_kl = old_per_token_logps - ref_per_token_logps + # Apply sequence-level KL penalty to rewards (sum KL across tokens first, then apply to each sequence) + kl = (per_token_kl * completion_mask).sum(-1) + kl = gather(kl) # rewards are gathered, so kl must be too + rewards = rewards - self.beta * kl + + grouped_rewards = rewards.view(-1, num_generations) + mean_grouped_rewards = torch.nanmean(grouped_rewards, dim=1) + if num_generations > 1: + std_rewards = nanstd(grouped_rewards, dim=1) + else: # doesn't occur during training, but could occur in eval when num_generations_eval=1 + std_rewards = torch.zeros_like(mean_grouped_rewards) + + # RLOO advantages computation. The leave-one-out baseline averages over scorable siblings only: nansum drops + # unscorable rewards and the divisor is (scorable count − 1). A group with a single scorable completion yields + # 0/0 = NaN, and unscorable rows stay NaN; both are zeroed by nan_to_num below. + scorable_counts = (~torch.isnan(grouped_rewards)).sum(dim=1, keepdim=True) # (num_prompts, 1) + grouped_sum = torch.nansum(grouped_rewards, dim=1, keepdim=True) # (num_prompts, 1) + if num_generations > 1: + baselines = (grouped_sum - grouped_rewards) / (scorable_counts - 1) # (num_prompts, num_generations) + baselines = baselines.view(-1) # Flatten back to match rewards shape + advantages = rewards - baselines + else: # this case doesn't occur during training, but could in eval when num_generations_eval=1 + advantages = torch.zeros_like(rewards) + + # Normalize advantages over the scorable subset only (unscorable advantages are still NaN here). + if self.normalize_advantages: + advantages = (advantages - torch.nanmean(advantages)) / (nanstd(advantages) + 1e-4) + + # Unscorable completions carry no learning signal: zero their advantage to keep them from moving the policy. + advantages = torch.nan_to_num(advantages, nan=0.0) + + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) # for logging + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Calculate and log the mean KL divergence between current and reference model + if self.beta != 0.0: + mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_func_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) + rewards = (rewards_per_func * self.reward_weights.to(rewards_per_func.device).unsqueeze(0)).nansum(dim=1) + rewards[unscorable_mask] = torch.nan # exclude unscorable rows from the logged reward stats + self._metrics[mode]["reward"].append(torch.nanmean(rewards).item()) + self._metrics[mode]["reward_std"].append(nanstd(rewards).item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._logs["prompt"].extend(gather_object(prompts_text)) + self._logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._logs["advantages"].extend(all_process_advantages.tolist()) + + # Flush user-logged extra columns (from log_extra), gathering across processes. + # Keys must be sorted so that all ranks call gather_object in the same order, otherwise values + # get mis-attributed across columns (dict insertion order may differ between processes). + for column in sorted(self._pending_extra_logs): + self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column])) + self._pending_extra_logs.clear() + + # Flush user-logged metrics (from log_metric), averaging across processes. + # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values + # get mis-attributed across metrics (dict insertion order may differ between processes). + for name in sorted(self._pending_metrics): + values = self._pending_metrics[name] + local_mean = sum(values) / len(values) + global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item() + self._metrics[mode][name].append(global_mean) + self._pending_metrics.clear() + + if images is not None: + self._logs["images"].extend(gather_object(images)) + + output = { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "old_logps": old_logps, + "advantages": advantages, + } + if "pixel_values" in forward_kwargs: + output["pixel_values"] = forward_kwargs["pixel_values"] + if "image_grid_thw" in forward_kwargs: + output["image_grid_thw"] = forward_kwargs["image_grid_thw"] + if "pixel_attention_mask" in forward_kwargs: + output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] + if "spatial_shapes" in forward_kwargs: + output["spatial_shapes"] = forward_kwargs["spatial_shapes"] + if "image_sizes" in forward_kwargs: + output["image_sizes"] = forward_kwargs["image_sizes"] + if "token_type_ids" in forward_kwargs: + output["token_type_ids"] = forward_kwargs["token_type_ids"] + if "mm_token_type_ids" in forward_kwargs: + output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"] + if "image_position_ids" in forward_kwargs: + output["image_position_ids"] = forward_kwargs["image_position_ids"] + if images is not None: + output["num_images"] = num_images + if num_tiles is not None: + output["num_tiles"] = num_tiles + return output + + @profiling_decorator + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The RLOOTrainer does not support returning outputs") + return self._compute_loss(model, inputs) + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # Compute the per_token_logps and the entropy at each position in the completion + per_token_logps, entropies, aux_loss = self._get_per_token_logps_and_entropies( + model, + input_ids, + attention_mask, + logits_to_keep, + compute_entropy=True, + compute_aux_loss=self.aux_loss_enabled, + pixel_values=inputs.get("pixel_values"), + image_grid_thw=inputs.get("image_grid_thw"), + num_images=inputs.get("num_images"), + pixel_attention_mask=inputs.get("pixel_attention_mask"), + spatial_shapes=inputs.get("spatial_shapes"), + num_tiles=inputs.get("num_tiles"), + image_sizes=inputs.get("image_sizes"), + token_type_ids=inputs.get("token_type_ids"), + mm_token_type_ids=inputs.get("mm_token_type_ids"), + image_position_ids=inputs.get("image_position_ids"), + ) + + logps = (per_token_logps * completion_mask).sum(1) # mask out padding and tokens after EOS + old_logps = inputs["old_logps"] + log_ratio = logps - old_logps + + # Compute the loss + advantages = inputs["advantages"] + coef_1 = torch.exp(log_ratio) + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + per_sequence_loss1 = coef_1 * advantages + per_sequence_loss2 = coef_2 * advantages + per_sequence_loss = -torch.min(per_sequence_loss1, per_sequence_loss2) + loss = per_sequence_loss.mean() + + # Log the metrics + mode = "train" if self.model.training else "eval" + + # RLOO returns an unscaled loss (the HF Trainer divides by gradient accumulation), so add the aux term unscaled + if self.aux_loss_enabled: + loss = loss + self.router_aux_loss_coef * aux_loss + self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) + + # Entropy + mean_entropy = (entropies * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) + + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) + is_region_clipped = is_low_clipped | is_high_clipped + gathered_low_clip = self.accelerator.gather(is_low_clipped.float().mean()) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(is_high_clipped.float().mean()) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(is_region_clipped.float().mean()) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + return loss + + # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and + # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + # Average the metrics + metrics = {} + for key, val in self._metrics[mode].items(): + # Filter out NaN values before averaging. A reward function that returns None for all samples + # in a batch produces NaN for that batch's metric. With logging_steps > 1, a naive sum()/len() + # would let a single NaN contaminate valid data from other batches. Only return None when no + # valid values remain (e.g. JSON loggers crash on float NaN). + valid = [v for v in val if not math.isnan(v)] + metrics[key] = sum(valid) / len(valid) if valid else None + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._logs["prompt"], + self._logs["completion"], + self._logs["rewards"], + self._logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + extra=dict(self._logs["extra"]), + ) + + logging_backends = [] + if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: + logging_backends.append(wandb) + if self.args.report_to and "trackio" in self.args.report_to: + logging_backends.append(trackio) + + table = { + "step": [self.state.global_step] * len(self._logs["prompt"]), + "prompt": self._logs["prompt"], + "completion": self._logs["completion"], + **self._logs["rewards"], + **self._logs["extra"], + "advantage": self._logs["advantages"], + } + + df_base = pd.DataFrame(table) + images_raw = self._logs["images"] or [] + + for logging_backend in logging_backends: + if images_raw: + images = [] + for image_list in self._logs["images"]: + images.append([logging_backend.Image(image) for image in image_list]) + df = pd.concat( + [df_base, pd.Series(images, name="image")], + axis=1, + copy=False, + ) + else: + df = df_base + + if self.log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + + logging_backend.log({"completions": logging_backend.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) +class UnslothRLOOTrainer(_UnslothRLOOTrainer): + """ + + Trainer for the Reinforce Leave One Out (RLOO) method. This algorithm was initially proposed in the paper [Back to + Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in + LLMs](https://huggingface.co/papers/2402.14740). + + Example: + + ```python + >>> from trl import RLOOTrainer + >>> from trl.rewards import accuracy_reward + >>> from datasets import load_dataset + + >>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train") + + >>> trainer = RLOOTrainer( + ... model="Qwen/Qwen2.5-0.5B-Instruct", + ... reward_funcs=accuracy_reward, + ... train_dataset=dataset, + ... ) + >>> trainer.train() + ``` + + Args: + model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `.from_pretrained` (where `` is derived from the model + config) with the keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + - A [`~peft.PeftModel`] object. Only causal language models are supported. + reward_funcs (`RewardFunc | list[RewardFunc]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can be either synchronous or asynchronous and can also return `None` when the reward is + not applicable to those samples. This is useful for multi-task training where different reward + functions apply to different types of samples. When a reward function returns `None` for a sample, + that reward function is excluded from the reward calculation for that sample. For more details, see + [Using a custom reward + function](#using-a-custom-reward-function). + + The trainer's state is also passed to the reward function. The trainer's state is an instance of + [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the + reward function's signature. + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + args ([`RLOOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + + """ + def __init__( + self, + model, + reward_funcs, + args = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + reward_processing_classes = None, + callbacks = None, + peft_config = None, + **kwargs + ): + if args is None: args = UnslothRLOOConfig() + use_bf16 = getattr(args, 'bf16', False) + if type(use_bf16) is not bool: use_bf16 = False + use_fp16 = getattr(args, 'fp16', False) + if type(use_fp16) is not bool: use_fp16 = False + force_float32 = False + full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' + if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): + print('Unsloth: Switching to float32 training since model cannot work with float16') + force_float32 = True + mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') + dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) + if dtype is None: dtype = model.get_input_embeddings().weight.dtype + from unsloth_zoo.utils import _get_dtype + dtype = _get_dtype(dtype) + float16 = dtype == torch.float16 + if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') + if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') + if force_float32: + # Forced float32 training + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': + # Mixed precision training + args.fp16 = float16 + args.bf16 = not float16 + os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' + # args.mixed_precision is a new argument which needs to be set now + elif mixed_precision_dtype == 'bfloat16': + # Both False since bfloat16 full finetuning doesn't do any autocasting. + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + + if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': + args.eval_strategy = 'steps' + if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 + ga_steps = getattr(args, 'gradient_accumulation_steps', None) + if ga_steps is not None and ga_steps > 1: + from transformers import __version__ as transformers_version + if Version(transformers_version) <= Version('4.45.2'): + print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' + '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') + if getattr(args, 'eval_strategy', 'no') != 'no': + eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) + if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size + if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps + fp16_full_eval = getattr(args, 'fp16_full_eval', False) + if type(fp16_full_eval) is not bool: fp16_full_eval = False + bf16_full_eval = getattr(args, 'bf16_full_eval', False) + if type(bf16_full_eval) is not bool: bf16_full_eval = False + if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True + if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False + if force_float32: + args.bf16_full_eval = False + args.fp16_full_eval = False + elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': + args.bf16_full_eval = True + args.fp16_full_eval = False + elif not bf16_full_eval and not fp16_full_eval: + args.bf16_full_eval = args.bf16 + args.fp16_full_eval = args.fp16 + _output_logits = False + if locals().get('compute_metrics', None) is not None: _output_logits = True + if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True + if _output_logits: + os.environ['UNSLOTH_RETURN_LOGITS'] = '1' + if model is not None: + _warnings_issued = getattr(model, 'warnings_issued', None) + if _warnings_issued is None: + model.warnings_issued = {} + elif not isinstance(_warnings_issued, dict): + try: + model.warnings_issued = dict(_warnings_issued) + except Exception: + model.warnings_issued = {} + if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): + pass + else: + model_max_seq_length = getattr(model, 'max_seq_length', None) + args_max_seq_length = getattr(args, 'max_seq_length', None) + if args_max_seq_length is None and model_max_seq_length is not None: + max_seq_length = model.max_seq_length + if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length + elif args_max_seq_length is not None and model_max_seq_length is not None: + if args_max_seq_length > model_max_seq_length: + print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' + 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') + args.max_seq_length = model_max_seq_length + if model is not None and hasattr(model, 'for_training'): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' + if 'processing_class' in locals(): + if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' + if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('rloo_trainer', other_metrics) + + # [TODO] Fix up DataParallel multiplying batch sizes + # [TODO] DDP works, but DP seems to not work? [TODO] + if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: + if getattr(args, "_n_gpu", 1) != 1: + args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + super().__init__( + model = model, + reward_funcs = reward_funcs, + args = args, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + reward_processing_classes = reward_processing_classes, + callbacks = callbacks, + peft_config = peft_config,**kwargs) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() + if hasattr(self, 'neftune_hook_handle'): + self.neftune_hook_handle.remove() + if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle + if getattr(args, 'neftune_noise_alpha', None) is not None: + model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha + pass + if hasattr(self, 'accelerator'): + scaler = self.accelerator.scaler + current_model = model + while hasattr(current_model, 'model'): + current_model.accelerator_scaler = scaler + current_model = current_model.model + current_model.accelerator_scaler = scaler + pass + if hasattr(self, 'train'): + self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) + pass + if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): + _vllm_tok = self.llm.get_tokenizer() + _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) + if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: + _vllm_tok.chat_template = _pc.chat_template + pass + +pass + + +if hasattr(logger, "addFilter"): + import logging + class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not (self.text in x.getMessage()) + pass + logger.addFilter(HideLoggingMessage("`use_cache=True`")) + diff --git a/unsloth_compiled_cache/UnslothRewardTrainer.py b/unsloth_compiled_cache/UnslothRewardTrainer.py new file mode 100644 index 0000000000000000000000000000000000000000..3c507835bf034ac5b60ee8241cd43bc2d2580f58 --- /dev/null +++ b/unsloth_compiled_cache/UnslothRewardTrainer.py @@ -0,0 +1,1480 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.reward_trainer import (Any, AutoModelForSequenceClassification, AutoTokenizer, Callable, DataCollator, DataCollatorForPreference, Dataset, EvalPrediction, IterableDataset, PartialState, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, RewardConfig, RewardTrainer, TrainerCallback, Version, _BaseTrainer, clone_chat_template, contextlib, create_model_from_path, dataclass, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_dataset_column_names, get_peft_model, is_conversational, is_peft_available, is_peft_model, json, logger, logging, nn, os, pad, peft, re, set_seed, suppress_seqcls_warning, torch, transformers, warnings, Any, AutoModelForSequenceClassification, AutoTokenizer, Callable, DataCollator, DataCollatorForPreference, Dataset, EvalPrediction, IterableDataset, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, RewardConfig, TrainerCallback, Version, clone_chat_template, contextlib, create_model_from_path, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, is_peft_available, is_peft_model, logger, os, pad, peft, re, set_seed, suppress_seqcls_warning, torch, transformers, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, re, torch) + + +import os +import math +import logging +from typing import * +from dataclasses import dataclass, field +from packaging.version import Version +import torch +import numpy as np +from contextlib import nullcontext +from torch.nn import functional as F +import inspect +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling +from transformers.training_args import ParallelMode +from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize + +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass +# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain +# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. +try: + from unsloth.models._utils import _unsloth_reset_stray_compile_cache +except Exception: + def _unsloth_reset_stray_compile_cache(self): pass +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Drop any torch.compile graph cache poisoned by a stray pre-train forward. + try: + _unsloth_reset_stray_compile_cache(self) + except Exception: + pass + # Finish the previous W&B run if this is a subsequent train() call. + # We do this at the START of train() (not the end) so that + # evaluate() / log() still work after train() completes. + # HF's WandbCallback.setup() will call wandb.init() for the new run. + # See: https://github.com/unslothai/unsloth/issues/3954 + if getattr(self, '_unsloth_training_completed', False): + try: + import wandb + if wandb.run is not None: + wandb.finish() + # Reset HF's WandbCallback so it calls wandb.init() for the new run + for cb in self.callback_handler.callbacks: + if type(cb).__name__ == 'WandbCallback': + cb._initialized = False + break + except: + pass + # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + output = f(self, *args, **kwargs) + # Restore previous mode when possible + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass + # Mark that training completed so the next train() call can + # finish this W&B run before starting a new one + self._unsloth_training_completed = True + return output + return wrapper +pass + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_hidden_states_selective_log_softmax( + hidden_states: torch.Tensor, + lm_head: torch.Tensor, + index: torch.Tensor, + chunks: int = 4, + logit_scale_multiply: float = 0.0, + logit_scale_divide: float = 0.0, + logit_softcapping: float = 0.0, + temperature: float = 1.0, +) -> torch.Tensor: + # All Unsloth Zoo code licensed under AGPL3 + flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + flat_index = index.reshape(-1) + + chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) + chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) + + all_per_token_logps = [] + + for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): + chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() + + if logit_scale_multiply != 0.0: + chunk_logits = chunk_logits * logit_scale_multiply + if logit_scale_divide != 0.0: + chunk_logits = chunk_logits / logit_scale_divide + if logit_softcapping != 0.0: + chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) + + chunk_logits = chunk_logits.to(torch.float32) + + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + + selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + + all_per_token_logps = torch.concat(all_per_token_logps) + + all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) + return all_per_token_logps + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_selective_log_softmax( + logits, + index, + temperature: float = 1.0, + chunks: int = 4, +): + chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) + chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) + all_per_token_logps = [] + # Per-chunk selective_log_softmax. + for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): + chunk_logits = chunk_logits.to(torch.float32) + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + pass + all_per_token_logps = torch.concat(all_per_token_logps) + all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) + return all_per_token_logps + +def calculate_pad_tokens_in_prompt( + input_ids: torch.Tensor, + logits_to_keep: int, + pad_token_id: int +) -> torch.Tensor: + """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" + if logits_to_keep >= input_ids.shape[1]: + raise ValueError("logits_to_keep must be smaller than the sequence length.") + + prompt_section = input_ids[:, :-logits_to_keep] + + padding_mask = (prompt_section == pad_token_id) + + pad_token_counts = padding_mask.sum(dim=1) + + return pad_token_counts + +def create_completion_attention_mask( + completion_input_ids: torch.Tensor, + left_pad_tokens_per_prompt: torch.Tensor, + max_left_pad: int, + pad_token_id: int +) -> torch.Tensor: + """Build a completion mask that zeros leading prompt and trailing pad tokens. + + For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) + this returns [0,0,0,1,1,1,0,0,0]. + """ + batch_size, completion_len = completion_input_ids.shape + device = completion_input_ids.device + + num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt + + indices = torch.arange(completion_len, device=device).unsqueeze(0) + shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) + + non_padding_mask = (completion_input_ids != pad_token_id) + + final_mask = shift_mask & non_padding_mask + + return final_mask + +def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: + """Move all padding tokens in each sequence to the right.""" + mask = (tensor != pad_id) + # stable=True since the binary mask is unordered. + sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) + packed_tensor = torch.gather(tensor, 1, sorted_indices) + return packed_tensor + +def align_logprobs_with_mask( + logprob_tensor: torch.Tensor, + attention_mask: torch.Tensor, + pad_value: float = 0.0 +) -> torch.Tensor: + """Align a log probability tensor with a given attention mask.""" + + device = logprob_tensor.device + batch_size, logprob_seq_len = logprob_tensor.shape + mask_seq_len = attention_mask.shape[1] + + padded_logprobs = torch.full( + attention_mask.shape, + fill_value=pad_value, + dtype=logprob_tensor.dtype, + device=device + ) + + left_pad_counts = torch.argmax(attention_mask, dim=1) + + cols = torch.arange(logprob_seq_len, device=device) + dest_indices = left_pad_counts.unsqueeze(1) + cols + + # Destination row indices, shape [batch_size, logprob_seq_len]. + row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) + + # Keep only in-bounds destinations, then scatter via advanced indexing. + valid_mask = dest_indices < mask_seq_len + valid_rows = row_indices[valid_mask] + valid_cols = dest_indices[valid_mask] + valid_vals = logprob_tensor[valid_mask] + padded_logprobs[valid_rows, valid_cols] = valid_vals + + return padded_logprobs + +def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" + if tool_mask is None: + return completion_mask + if tool_mask.shape[0] != completion_mask.shape[0]: + raise ValueError("tool_mask batch size must match completion_mask batch size.") + + tool_mask = tool_mask.to(device=completion_mask.device) + if tool_mask.shape == completion_mask.shape: + aligned_tool_mask = tool_mask + else: + aligned_tool_mask = align_logprobs_with_mask( + tool_mask, + completion_mask, + pad_value=0, + ) + return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) + +def autotune_batch_and_chunks( + total_input_rows, + seq_len, + hidden_size, + vocab_size, + dtype_bytes=16, + multiplier=None +): + if multiplier is None: + final_m = max(4, seq_len // 4096) + else: + final_m = multiplier + + if torch.cuda.is_available(): + free_bytes, _ = torch.cuda.mem_get_info() + limit_gb = (free_bytes / (1024**3))*.80 + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + # XPU: estimate free memory as total - reserved. + total_mem = torch.xpu.get_device_properties(0).total_memory + reserved_mem = torch.xpu.memory_reserved() + free_bytes = total_mem - reserved_mem + limit_gb = (free_bytes / (1024**3)) * 0.80 + else: + # Fallback: assume 8GB available. + limit_gb = 8.0 + + bytes_to_gb = 1024**3 + + b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) + + hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb + + base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb + logits_gb = base_logits / final_m + + total_mem_gb = hidden_gb + logits_gb + + valid_mask = total_mem_gb <= limit_gb + valid_indices = torch.nonzero(valid_mask, as_tuple=False) + + if valid_indices.shape[0] == 0: + #This means your GPU will OOM + return 4, final_m + + best_idx = valid_indices[0].item() + final_b = int(b_vals[best_idx].item()) + + return final_b, final_m + +def sanitize_logprob(logprob): + """Local port of trl.scripts.vllm_serve.sanitize_logprob. + Filters NaN logprobs from vLLM outputs.""" + value = logprob.logprob + if math.isnan(value): + logging.getLogger(__name__).warning( + f"Generated NaN logprob, token logprob '{logprob}' will be ignored" + ) + return None + return value +@dataclass +class UnslothRewardConfig(RewardConfig): + """ + + Configuration class for the [`RewardTrainer`]. + + This class includes only the parameters that are specific to Reward training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`RewardTrainer`] is provided as a string. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to + [`~transformers.AutoModelForSequenceClassification.from_pretrained`] and + [`~transformers.AutoTokenizer.from_pretrained`]. + chat_template_path (`str`, *optional*): + If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory + or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must + ensure that any special tokens referenced in the template are added to the tokenizer and that the model's + embedding layer is resized accordingly. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + + > Parameters that control the data preprocessing + + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + eos_token (`str`, *optional*): + Token used to indicate the end of a turn or sequence. If `None`, it defaults to + `processing_class.eos_token`. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Samples are filtered out if either chosen or rejected sequence + exceeds this value. If `None`, no filtering is applied. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + + > Parameters that control the training + + center_rewards_coefficient (`float`, *optional*): + Coefficient to incentivize the reward model to output mean-zero rewards (proposed by + https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`. + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + + > Deprecated parameters + + pad_token: + + + + Parameter `pad_token` is deprecated and will be removed in version v2.0.0. Set `tokenizer.pad_token` + directly and pass it as `processing_class` to the trainer instead. + + + + > [!NOTE] + > These parameters have default values different from [`~transformers.TrainingArguments`]: + > - `logging_steps`: Defaults to `10` instead of `500`. + > - `gradient_checkpointing`: Defaults to `True` instead of `False`. + > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. + > - `learning_rate`: Defaults to `1e-4` instead of `5e-5`. + + """ + vllm_sampling_params: Optional[Any] = field( + default = None, + metadata = {'help': 'vLLM SamplingParams'}, + ) + unsloth_num_chunks : Optional[int] = field( + default = -1, + metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, + ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {'help': 'Multiplier for chunked logit computations.'}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, + ) + max_seq_length : Optional[int] = field( + default = None, + metadata = {'help': 'Maximum sequence length to truncate to.'}, + ) + def __init__( + self, + output_dir = None, + per_device_train_batch_size = 4, + num_train_epochs = 3.0, + max_steps = -1, + learning_rate = 5e-05, + lr_scheduler_type = 'linear', + lr_scheduler_kwargs = None, + warmup_steps = 0.1, + optim = 'adamw_8bit', + optim_args = None, + weight_decay = 0.001, + adam_beta1 = 0.9, + adam_beta2 = 0.999, + adam_epsilon = 1e-08, + optim_target_modules = None, + gradient_accumulation_steps = 2, + average_tokens_across_devices = True, + max_grad_norm = 1.0, + label_smoothing_factor = 0.0, + bf16 = False, + fp16 = False, + bf16_full_eval = False, + fp16_full_eval = False, + tf32 = None, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = None, + torch_compile = False, + torch_compile_backend = None, + torch_compile_mode = None, + use_liger_kernel = False, + liger_kernel_config = None, + use_cache = False, + neftune_noise_alpha = None, + torch_empty_cache_steps = 250, + auto_find_batch_size = False, + logging_strategy = 'steps', + logging_steps = 1, + logging_first_step = False, + log_on_each_node = True, + logging_nan_inf_filter = False, + include_num_input_tokens_seen = False, + log_level = 'passive', + log_level_replica = 'warning', + disable_tqdm = None, + report_to = 'none', + run_name = None, + project = 'huggingface', + trackio_space_id = 'trackio', + eval_strategy = 'no', + eval_steps = None, + eval_delay = 0, + per_device_eval_batch_size = 4, + prediction_loss_only = False, + eval_on_start = False, + eval_do_concat_batches = True, + eval_use_gather_object = False, + eval_accumulation_steps = 2, + batch_eval_metrics = False, + save_only_model = False, + save_strategy = 'steps', + save_steps = 500, + save_on_each_node = False, + save_total_limit = None, + enable_jit_checkpoint = False, + push_to_hub = False, + hub_token = None, + hub_private_repo = None, + hub_model_id = None, + hub_strategy = 'every_save', + hub_always_push = False, + hub_revision = None, + load_best_model_at_end = False, + metric_for_best_model = None, + greater_is_better = None, + ignore_data_skip = False, + restore_callback_states_from_checkpoint = False, + full_determinism = False, + seed = 3407, + data_seed = 3407, + use_cpu = False, + accelerator_config = None, + parallelism_config = None, + dataloader_drop_last = False, + dataloader_num_workers = 0, + dataloader_pin_memory = True, + dataloader_persistent_workers = False, + dataloader_prefetch_factor = None, + remove_unused_columns = True, + label_names = None, + train_sampling_strategy = 'random', + length_column_name = 'length', + ddp_find_unused_parameters = None, + ddp_bucket_cap_mb = None, + ddp_broadcast_buffers = None, + ddp_backend = None, + ddp_timeout = 1800, + fsdp = None, + fsdp_config = None, + deepspeed = None, + debug = '', + skip_memory_metrics = True, + do_train = False, + do_eval = False, + do_predict = False, + resume_from_checkpoint = None, + warmup_ratio = None, + logging_dir = None, + local_rank = -1, + model_init_kwargs = None, + trust_remote_code = False, + chat_template_path = None, + disable_dropout = True, + dataset_num_proc = None, + eos_token = None, + max_length = 1024, + pad_to_multiple_of = None, + center_rewards_coefficient = None, + activation_offloading = False, + pad_token = None, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, + max_seq_length = None, + **kwargs, + ): + if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') + if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') + if num_train_epochs is None: + num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override + if output_dir is None and save_strategy == 'steps' and save_steps == 500: + output_dir = 'unsloth_training_checkpoints' + save_strategy = 'no' + import multiprocessing as _mp + if dataset_num_proc is None: + if _mp.get_start_method() != 'fork': + dataset_num_proc = None + else: + import psutil + dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64) + memory_gb_left = psutil.virtual_memory().available / (1024**3) + if memory_gb_left <= 2: dataset_num_proc = 1 + else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left)) + if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': + from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION + if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: + from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE + pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE + + + super().__init__( + output_dir = output_dir, + per_device_train_batch_size = per_device_train_batch_size, + num_train_epochs = num_train_epochs, + max_steps = max_steps, + learning_rate = learning_rate, + lr_scheduler_type = lr_scheduler_type, + lr_scheduler_kwargs = lr_scheduler_kwargs, + warmup_steps = warmup_steps, + optim = optim, + optim_args = optim_args, + weight_decay = weight_decay, + adam_beta1 = adam_beta1, + adam_beta2 = adam_beta2, + adam_epsilon = adam_epsilon, + optim_target_modules = optim_target_modules, + gradient_accumulation_steps = gradient_accumulation_steps, + average_tokens_across_devices = average_tokens_across_devices, + max_grad_norm = max_grad_norm, + label_smoothing_factor = label_smoothing_factor, + bf16 = bf16, + fp16 = fp16, + bf16_full_eval = bf16_full_eval, + fp16_full_eval = fp16_full_eval, + tf32 = tf32, + gradient_checkpointing = gradient_checkpointing, + gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, + torch_compile = torch_compile, + torch_compile_backend = torch_compile_backend, + torch_compile_mode = torch_compile_mode, + use_liger_kernel = use_liger_kernel, + liger_kernel_config = liger_kernel_config, + use_cache = use_cache, + neftune_noise_alpha = neftune_noise_alpha, + torch_empty_cache_steps = torch_empty_cache_steps, + auto_find_batch_size = auto_find_batch_size, + logging_strategy = logging_strategy, + logging_steps = logging_steps, + logging_first_step = logging_first_step, + log_on_each_node = log_on_each_node, + logging_nan_inf_filter = logging_nan_inf_filter, + include_num_input_tokens_seen = include_num_input_tokens_seen, + log_level = log_level, + log_level_replica = log_level_replica, + disable_tqdm = disable_tqdm, + report_to = report_to, + run_name = run_name, + project = project, + trackio_space_id = trackio_space_id, + eval_strategy = eval_strategy, + eval_steps = eval_steps, + eval_delay = eval_delay, + per_device_eval_batch_size = per_device_eval_batch_size, + prediction_loss_only = prediction_loss_only, + eval_on_start = eval_on_start, + eval_do_concat_batches = eval_do_concat_batches, + eval_use_gather_object = eval_use_gather_object, + eval_accumulation_steps = eval_accumulation_steps, + batch_eval_metrics = batch_eval_metrics, + save_only_model = save_only_model, + save_strategy = save_strategy, + save_steps = save_steps, + save_on_each_node = save_on_each_node, + save_total_limit = save_total_limit, + enable_jit_checkpoint = enable_jit_checkpoint, + push_to_hub = push_to_hub, + hub_token = hub_token, + hub_private_repo = hub_private_repo, + hub_model_id = hub_model_id, + hub_strategy = hub_strategy, + hub_always_push = hub_always_push, + hub_revision = hub_revision, + load_best_model_at_end = load_best_model_at_end, + metric_for_best_model = metric_for_best_model, + greater_is_better = greater_is_better, + ignore_data_skip = ignore_data_skip, + restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, + full_determinism = full_determinism, + seed = seed, + data_seed = data_seed, + use_cpu = use_cpu, + accelerator_config = accelerator_config, + parallelism_config = parallelism_config, + dataloader_drop_last = dataloader_drop_last, + dataloader_num_workers = dataloader_num_workers, + dataloader_pin_memory = dataloader_pin_memory, + dataloader_persistent_workers = dataloader_persistent_workers, + dataloader_prefetch_factor = dataloader_prefetch_factor, + remove_unused_columns = remove_unused_columns, + label_names = label_names, + train_sampling_strategy = train_sampling_strategy, + length_column_name = length_column_name, + ddp_find_unused_parameters = ddp_find_unused_parameters, + ddp_bucket_cap_mb = ddp_bucket_cap_mb, + ddp_broadcast_buffers = ddp_broadcast_buffers, + ddp_backend = ddp_backend, + ddp_timeout = ddp_timeout, + fsdp = fsdp, + fsdp_config = fsdp_config, + deepspeed = deepspeed, + debug = debug, + skip_memory_metrics = skip_memory_metrics, + do_train = do_train, + do_eval = do_eval, + do_predict = do_predict, + resume_from_checkpoint = resume_from_checkpoint, + warmup_ratio = warmup_ratio, + logging_dir = logging_dir, + local_rank = local_rank, + model_init_kwargs = model_init_kwargs, + trust_remote_code = trust_remote_code, + chat_template_path = chat_template_path, + disable_dropout = disable_dropout, + dataset_num_proc = dataset_num_proc, + eos_token = eos_token, + max_length = max_length, + pad_to_multiple_of = pad_to_multiple_of, + center_rewards_coefficient = center_rewards_coefficient, + activation_offloading = activation_offloading, + pad_token = pad_token,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier + self.max_seq_length = max_seq_length + # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+ + if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: + if 'use_reentrant' in self.gradient_checkpointing_kwargs: + del self.gradient_checkpointing_kwargs['use_reentrant'] + +pass + +class _UnslothRewardTrainer(_BaseTrainer): + """""" + + _tag_names = ["trl", "reward-trainer"] + _name = "Reward" + _template_file = "rm_model_card.md" + + def __init__( + self, + model: "str | PreTrainedModel | PeftModel", + args: RewardConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: "PeftConfig | None" = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = RewardConfig(f"{model_name}-Reward") + + if train_dataset is None: + raise ValueError("`train_dataset` is required") + elif isinstance(train_dataset, IterableDataset): + # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate + # batches from multiple processes, leading to mismatch errors. + if args.accelerator_config.dispatch_batches is True: + logger.warning( + "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " + "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " + "`dispatch_batches` in `RewardConfig` or set it to `False`." + ) + args.accelerator_config.dispatch_batches = False + + # Model + # As AutoModelForSequenceClassification.from_pretrained[] will add a random head for the model, set_seed must + # be done before loading the model to ensure reproducibility. + set_seed(args.seed) + if isinstance(model, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs["num_labels"] = 1 # the only output of the model is the reward score + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + with suppress_seqcls_warning(): + model = create_model_from_path(model, AutoModelForSequenceClassification, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `RewardConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + # Validate that the model has num_labels = 1 [required for reward models] + if getattr(model.config, "num_labels", None) != 1: + raise ValueError( + f"The model has `num_labels={model.config.num_labels}`, but reward models require `num_labels=1` " + "to output a single scalar reward per sequence. Please instantiate your model with `num_labels=1` " + "or pass a model name as a string to have it configured automatically." + ) + # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do + _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + + # Processing class + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained( + get_config_model_id(model.config), trust_remote_code=args.trust_remote_code + ) + + # Handle pad token for processors or tokenizers + if args.eos_token is not None: + if args.eos_token not in processing_class.get_vocab(): + raise ValueError( + f"The specified `eos_token` ('{args.eos_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " + "in the vocabulary before using it as an EOS token." + ) + processing_class.eos_token = args.eos_token + + if args.chat_template_path is not None: + if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): + with open(args.chat_template_path, encoding="utf-8") as chat_template_file: + processing_class.chat_template = chat_template_file.read() + added_tokens = [] + else: + model, processing_class, added_tokens = clone_chat_template( + model, processing_class, args.chat_template_path + ) + else: + added_tokens = [] + + # PEFT + if False: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if is_peft_model(model): + raise ValueError( + "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " + "and unload the existing adapter, save the resulting base model, and then pass that base model along " + "with the new `peft_config` to the trainer." + ) + if added_tokens: + # Ensure that the added tokens are trainable + if peft_config.trainable_token_indices is None: + peft_config.trainable_token_indices = {"embed_tokens": added_tokens} + elif "embed_tokens" not in peft_config.trainable_token_indices: + peft_config.trainable_token_indices["embed_tokens"] = added_tokens + else: + peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) + # Ensure that the lm_head is trainable + if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: + logger.warning( + "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " + "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " + "tokens, leading to degraded generation quality. To fix this, add " + "`modules_to_save=['lm_head']` to your PEFT configuration." + ) + + if peft_config.modules_to_save is None: + peft_config.modules_to_save = ["lm_head"] + else: + peft_config.modules_to_save.append("lm_head") + # Create PEFT model + # ZeRO-3 + PEFT for non-quantized models: + # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. + # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, + # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. + # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. + # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for + # non-quantized bf16 training. + # - See: + # - TRL issue: https://github.com/huggingface/trl/issues/6089 + # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 + # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg + get_peft_model_kwargs = {} + if ( + args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and not _is_quantized_model + and Version(peft.__version__) >= Version("0.12.0") + ): + get_peft_model_kwargs["autocast_adapter_dtype"] = False + model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + + # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally + # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 + if is_peft_model(model) and args.gradient_checkpointing: + model.enable_input_require_grads() + + # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the + # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by + # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for + # quantized models. See: https://github.com/huggingface/peft/issues/2889 + if _is_quantized_model: + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.bfloat16) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + # Pad token [needed for SequenceClassification models] + # If not provided, use the one from the processing class or the eos token if the processing class does not have + # a pad token. + pad_token = args.pad_token or processing_class.pad_token or processing_class.eos_token + if pad_token not in processing_class.get_vocab(): + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + processing_class.pad_token = pad_token + # SequenceClassification models need `config.pad_token_id` to locate the last non-pad token. + model.config.pad_token_id = processing_class.pad_token_id + + # Data collator + if data_collator is None: + data_collator = DataCollatorForPreference( + pad_token_id=processing_class.pad_token_id, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + + # Dataset + train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # During evaluation, Trainer calls compute_loss[] only if can_return_loss is True and label_names is empty. + self.can_return_loss = True + self.label_names = [] + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase, + args: RewardConfig, + dataset_name: str, + ) -> Dataset | IterableDataset: + # If the dataset is already preprocessed (tokenized), skip the processing steps. + column_names = get_dataset_column_names(dataset) + is_processed = "chosen_ids" in column_names and "rejected_ids" in column_names + has_legacy_processed_columns = "chosen_input_ids" in column_names and "rejected_input_ids" in column_names + if has_legacy_processed_columns and not is_processed: + warnings.warn( + "Detected legacy dataset columns `chosen_input_ids`/`rejected_input_ids`; they are deprecated and " + "will not be supported in v1. Please migrate to `chosen_ids`/`rejected_ids`.", + FutureWarning, + stacklevel=2, + ) + dataset = dataset.rename_column("chosen_input_ids", "chosen_ids") + dataset = dataset.rename_column("rejected_input_ids", "rejected_ids") + is_processed = True + + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + with PartialState().main_process_first(): + if not is_processed: + # Add EOS token if needed: non-conversational only + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if not example["chosen"].endswith(eos_token): + example["chosen"] = example["chosen"] + eos_token + if "rejected" in example and not example["rejected"].endswith(eos_token): + example["rejected"] = example["rejected"] + eos_token + return example + + dataset = dataset.map( + add_eos, + fn_kwargs={"eos_token": processing_class.eos_token}, + **map_kwargs, + ) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + def tokenize_fn(example, processing_class): + tools = example.get("tools") + tools = json.loads(tools) if isinstance(tools, str) else tools + if "prompt" in example: # explicit prompt case + example["chosen"] = example["prompt"] + example["chosen"] + example["rejected"] = example["prompt"] + example["rejected"] + + if is_conversational(example): + chosen_ids = processing_class.apply_chat_template( + example["chosen"], + tools=tools, + return_dict=True, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + rejected_ids = processing_class.apply_chat_template( + example["rejected"], + tools=tools, + return_dict=True, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + output = {"chosen_ids": chosen_ids, "rejected_ids": rejected_ids} + else: + output = { + "chosen_ids": processing_class(text=example["chosen"])["input_ids"], + "rejected_ids": processing_class(text=example["rejected"])["input_ids"], + } + return output + + dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) + + # Filter samples that are longer than `max_length` + if args.max_length is not None: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Filtering {dataset_name} >{args.max_length} tokens" + dataset = dataset.filter( + lambda example: len(example["chosen_ids"]) <= args.max_length + and len(example["rejected_ids"]) <= args.max_length, + **map_kwargs, + ) + + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). + if self._signature_columns is None: + self._signature_columns = ["chosen_ids", "rejected_ids", "margin"] + + def evaluate( + self, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> dict[str, float]: + # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way + # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer. `_prepare_dataset` is + # idempotent: it skips datasets that are already tokenized. A `str` selects a dataset that was already prepared + # at init time, so it's left untouched. + if eval_dataset is not None and not isinstance(eval_dataset, str): + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, self.processing_class, self.args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, self.processing_class, self.args, "eval") + return super().evaluate( + eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix + ) + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + mode = "train" if self.model.training else "eval" + + # If not set, defaults from model config and may warn since cache isn't compatible with gradient checkpointing + inputs["use_cache"] = False + outputs = model(**inputs) + + # Split the rewards into chosen and rejected + rewards_chosen, rewards_rejected = torch.chunk(outputs.logits.squeeze(-1), chunks=2) + + # Calculate loss, optionally modulate with margin + if "margin" in inputs: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean() + else: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() + + if self.args.center_rewards_coefficient is not None: + loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2) + + if mode == "train": + num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + self._total_train_tokens += num_tokens_in_batch + self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + + # Compute min, mean, max, accuracy and margin + with torch.no_grad(): + all_rewards = self.accelerator.gather(outputs.logits) + self._metrics[mode]["min_reward"].append(all_rewards.min().item()) + self._metrics[mode]["mean_reward"].append(all_rewards.mean().item()) + self._metrics[mode]["max_reward"].append(all_rewards.max().item()) + + mean_accuracy = (rewards_chosen > rewards_rejected).float().mean() + mean_accuracy = self.accelerator.gather_for_metrics(mean_accuracy).mean().item() + self._metrics[mode]["accuracy"].append(mean_accuracy) + + mean_margin = (rewards_chosen - rewards_rejected).mean() + mean_margin = self.accelerator.gather_for_metrics(mean_margin).mean() + self._metrics[mode]["margin"].append(mean_margin.item()) + + return (loss, outputs) if return_outputs else loss + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) +class UnslothRewardTrainer(_UnslothRewardTrainer): + """ + + Trainer for Outcome-supervised Reward Models (ORM). + + This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + >>> from trl import RewardTrainer + >>> from datasets import load_dataset + + >>> dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train") + + >>> trainer = RewardTrainer( + ... model="Qwen/Qwen2.5-0.5B-Instruct", + ... train_dataset=dataset, + ... ) + >>> trainer.train() + ``` + + Args: + model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `AutoModelForSequenceClassification.from_pretrained` with the keyword arguments in + `args.model_init_kwargs`. + - A sequence classification [`~transformers.PreTrainedModel`] object. + - A sequence classification [`~peft.PeftModel`] object. + args ([`RewardConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator ([`~transformers.DataCollator`], *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~trainer.reward_trainer.DataCollatorForPreference`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. This trainer supports [preference](#preference) type (both implicit and + explicit prompt). The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + + The trainer also supports processed datasets (tokenized) as long as they contain `chosen_ids` and + `rejected_ids` fields. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*): + Tokenizer used to process the data. If `None`, the tokenizer is loaded from the model's name with + [`~transformers.AutoTokenizer.from_pretrained`]. A padding token, `processing_class.pad_token`, must be + set. If the processing class has not set a padding token, `processing_class.eos_token` will be used as the + default. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a + [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing + [`RewardConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a + boolean `compute_result` argument. This will be triggered after the last eval batch to signal that the + function needs to calculate and return the global summary statistics rather than accumulating the + batch-level statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before + initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. Note that if the loaded + model is a causal LM, it's highly recommended to set `modules_to_save=["score"]` in the PEFT configuration + to ensure that the reward head is properly trained. + + """ + def __init__( + self, + model, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + compute_metrics = None, + callbacks = None, + optimizer_cls_and_kwargs = None, + preprocess_logits_for_metrics = None, + peft_config = None, + **kwargs + ): + if args is None: args = UnslothRewardConfig() + use_bf16 = getattr(args, 'bf16', False) + if type(use_bf16) is not bool: use_bf16 = False + use_fp16 = getattr(args, 'fp16', False) + if type(use_fp16) is not bool: use_fp16 = False + force_float32 = False + full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' + if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): + print('Unsloth: Switching to float32 training since model cannot work with float16') + force_float32 = True + mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') + dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) + if dtype is None: dtype = model.get_input_embeddings().weight.dtype + from unsloth_zoo.utils import _get_dtype + dtype = _get_dtype(dtype) + float16 = dtype == torch.float16 + if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') + if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') + if force_float32: + # Forced float32 training + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': + # Mixed precision training + args.fp16 = float16 + args.bf16 = not float16 + os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' + # args.mixed_precision is a new argument which needs to be set now + elif mixed_precision_dtype == 'bfloat16': + # Both False since bfloat16 full finetuning doesn't do any autocasting. + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + + if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': + args.eval_strategy = 'steps' + if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 + ga_steps = getattr(args, 'gradient_accumulation_steps', None) + if ga_steps is not None and ga_steps > 1: + from transformers import __version__ as transformers_version + if Version(transformers_version) <= Version('4.45.2'): + print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' + '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') + if getattr(args, 'eval_strategy', 'no') != 'no': + eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) + if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size + if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps + fp16_full_eval = getattr(args, 'fp16_full_eval', False) + if type(fp16_full_eval) is not bool: fp16_full_eval = False + bf16_full_eval = getattr(args, 'bf16_full_eval', False) + if type(bf16_full_eval) is not bool: bf16_full_eval = False + if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True + if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False + if force_float32: + args.bf16_full_eval = False + args.fp16_full_eval = False + elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': + args.bf16_full_eval = True + args.fp16_full_eval = False + elif not bf16_full_eval and not fp16_full_eval: + args.bf16_full_eval = args.bf16 + args.fp16_full_eval = args.fp16 + _output_logits = False + if locals().get('compute_metrics', None) is not None: _output_logits = True + if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True + if _output_logits: + os.environ['UNSLOTH_RETURN_LOGITS'] = '1' + if model is not None: + _warnings_issued = getattr(model, 'warnings_issued', None) + if _warnings_issued is None: + model.warnings_issued = {} + elif not isinstance(_warnings_issued, dict): + try: + model.warnings_issued = dict(_warnings_issued) + except Exception: + model.warnings_issued = {} + if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): + pass + else: + model_max_seq_length = getattr(model, 'max_seq_length', None) + args_max_seq_length = getattr(args, 'max_seq_length', None) + if args_max_seq_length is None and model_max_seq_length is not None: + max_seq_length = model.max_seq_length + if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length + elif args_max_seq_length is not None and model_max_seq_length is not None: + if args_max_seq_length > model_max_seq_length: + print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' + 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') + args.max_seq_length = model_max_seq_length + if model is not None and hasattr(model, 'for_training'): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' + if 'processing_class' in locals(): + if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' + if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' + __tokenizer = processing_class if 'processing_class' in locals() else tokenizer + from unsloth_zoo.vision_utils import UnslothVisionDataCollator + if not isinstance(data_collator, UnslothVisionDataCollator): + if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names: + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names: + data_collator = DataCollatorForSeq2Seq( + __tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + else: + if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False + if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' + if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} + if not isinstance(data_collator, UnslothVisionDataCollator): + if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): + if isinstance(data_collator, DataCollatorForSeq2Seq): + data_collator = DataCollatorForSeq2Seq( + __tokenizer.tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer.tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('reward_trainer', other_metrics) + + # [TODO] Fix up DataParallel multiplying batch sizes + # [TODO] DDP works, but DP seems to not work? [TODO] + if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: + if getattr(args, "_n_gpu", 1) != 1: + args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + super().__init__( + model = model, + args = args, + data_collator = data_collator, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + compute_metrics = compute_metrics, + callbacks = callbacks, + optimizer_cls_and_kwargs = optimizer_cls_and_kwargs, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config,**kwargs) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() + if hasattr(self, 'neftune_hook_handle'): + self.neftune_hook_handle.remove() + if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle + if getattr(args, 'neftune_noise_alpha', None) is not None: + model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha + pass + if hasattr(self, 'accelerator'): + scaler = self.accelerator.scaler + current_model = model + while hasattr(current_model, 'model'): + current_model.accelerator_scaler = scaler + current_model = current_model.model + current_model.accelerator_scaler = scaler + pass + if hasattr(self, 'train'): + self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) + pass + if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): + _vllm_tok = self.llm.get_tokenizer() + _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) + if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: + _vllm_tok.chat_template = _pc.chat_template + pass + +pass + + +if hasattr(logger, "addFilter"): + import logging + class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not (self.text in x.getMessage()) + pass + logger.addFilter(HideLoggingMessage("`use_cache=True`")) + diff --git a/unsloth_compiled_cache/UnslothSFTTrainer.py b/unsloth_compiled_cache/UnslothSFTTrainer.py new file mode 100644 index 0000000000000000000000000000000000000000..83a71a14fb2bef16c3f9075ebf73a8005e6229ea --- /dev/null +++ b/unsloth_compiled_cache/UnslothSFTTrainer.py @@ -0,0 +1,2006 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from trl.trainer.sft_trainer import (Any, AutoProcessor, Callable, DataCollator, DataCollatorForLanguageModeling, DataCollatorForVisionLanguageModeling, Dataset, EvalPrediction, F, FLASH_ATTENTION_VARIANTS, IterableDataset, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTConfig, SFTTrainer, TrainerCallback, TrainingArguments, Version, _BaseTrainer, _CHUNKED_LM_HEAD_CHUNK_SIZE, _chunk, _patch_chunked_ce_lm_head, apply_chat_template, clone_chat_template, contextlib, create_model_from_path, dataclass, defaultdict, dft_loss, get_act_offloading_ctx_manager, get_config_model_id, get_dataset_column_names, get_peft_model, get_training_chat_template, has_generation_markers, is_chat_template_stop_token_trained, is_conversational, is_peft_available, is_peft_model, logger, nn, os, pack_dataset, pad, peft, prepare_multimodal_messages, selective_log_softmax, torch, transformers, types, warnings, Any, AutoProcessor, Callable, DataCollator, DataCollatorForLanguageModeling, DataCollatorForVisionLanguageModeling, Dataset, EvalPrediction, F, FLASH_ATTENTION_VARIANTS, IterableDataset, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTConfig, SFTTrainer, TrainerCallback, TrainingArguments, Version, clone_chat_template, contextlib, create_model_from_path, defaultdict, dft_loss, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, get_training_chat_template, has_generation_markers, is_chat_template_stop_token_trained, is_conversational, is_peft_available, is_peft_model, logger, nn, os, pad, peft, torch, transformers, Callable, DataCollator, DataCollatorForLanguageModeling, Dataset, F, IterableDataset, apply_chat_template, is_conversational, os, pack_dataset, pad, transformers, warnings, F, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch, F, os) + + +import os +import math +import logging +from typing import * +from dataclasses import dataclass, field +from packaging.version import Version +import torch +import numpy as np +from contextlib import nullcontext +from torch.nn import functional as F +import inspect +from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling +from transformers.training_args import ParallelMode +from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize + +# Wrap trainer with padding to right and enable training mode +import functools +from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass +# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain +# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. +try: + from unsloth.models._utils import _unsloth_reset_stray_compile_cache +except Exception: + def _unsloth_reset_stray_compile_cache(self): pass +def prepare_for_training_mode(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + # Drop any torch.compile graph cache poisoned by a stray pre-train forward. + try: + _unsloth_reset_stray_compile_cache(self) + except Exception: + pass + # Finish the previous W&B run if this is a subsequent train() call. + # We do this at the START of train() (not the end) so that + # evaluate() / log() still work after train() completes. + # HF's WandbCallback.setup() will call wandb.init() for the new run. + # See: https://github.com/unslothai/unsloth/issues/3954 + if getattr(self, '_unsloth_training_completed', False): + try: + import wandb + if wandb.run is not None: + wandb.finish() + # Reset HF's WandbCallback so it calls wandb.init() for the new run + for cb in self.callback_handler.callbacks: + if type(cb).__name__ == 'WandbCallback': + cb._initialized = False + break + except: + pass + # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training + if hasattr(self, 'model') and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + output = f(self, *args, **kwargs) + # Restore previous mode when possible + if hasattr(self, 'model') and hasattr(self.model, "for_inference"): + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass + # Mark that training completed so the next train() call can + # finish this W&B run before starting a new one + self._unsloth_training_completed = True + return output + return wrapper +pass + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : False, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_hidden_states_selective_log_softmax( + hidden_states: torch.Tensor, + lm_head: torch.Tensor, + index: torch.Tensor, + chunks: int = 4, + logit_scale_multiply: float = 0.0, + logit_scale_divide: float = 0.0, + logit_softcapping: float = 0.0, + temperature: float = 1.0, +) -> torch.Tensor: + # All Unsloth Zoo code licensed under AGPL3 + flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + flat_index = index.reshape(-1) + + chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) + chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) + + all_per_token_logps = [] + + for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): + chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() + + if logit_scale_multiply != 0.0: + chunk_logits = chunk_logits * logit_scale_multiply + if logit_scale_divide != 0.0: + chunk_logits = chunk_logits / logit_scale_divide + if logit_softcapping != 0.0: + chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) + + chunk_logits = chunk_logits.to(torch.float32) + + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + + selected_logits = torch.gather(chunk_logits, dim=-1, index=chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim=-1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + + all_per_token_logps = torch.concat(all_per_token_logps) + + all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) + return all_per_token_logps + +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def chunked_selective_log_softmax( + logits, + index, + temperature: float = 1.0, + chunks: int = 4, +): + chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) + chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) + all_per_token_logps = [] + # Per-chunk selective_log_softmax. + for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): + chunk_logits = chunk_logits.to(torch.float32) + if temperature != 1.0: + chunk_logits = chunk_logits / temperature + selected_logits = torch.gather(chunk_logits, dim = -1, index = chunk_index.unsqueeze(-1)).squeeze(-1) + logsumexp_values = torch.logsumexp(chunk_logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values + all_per_token_logps.append(per_token_logps) + pass + all_per_token_logps = torch.concat(all_per_token_logps) + all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) + return all_per_token_logps + +def calculate_pad_tokens_in_prompt( + input_ids: torch.Tensor, + logits_to_keep: int, + pad_token_id: int +) -> torch.Tensor: + """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" + if logits_to_keep >= input_ids.shape[1]: + raise ValueError("logits_to_keep must be smaller than the sequence length.") + + prompt_section = input_ids[:, :-logits_to_keep] + + padding_mask = (prompt_section == pad_token_id) + + pad_token_counts = padding_mask.sum(dim=1) + + return pad_token_counts + +def create_completion_attention_mask( + completion_input_ids: torch.Tensor, + left_pad_tokens_per_prompt: torch.Tensor, + max_left_pad: int, + pad_token_id: int +) -> torch.Tensor: + """Build a completion mask that zeros leading prompt and trailing pad tokens. + + For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) + this returns [0,0,0,1,1,1,0,0,0]. + """ + batch_size, completion_len = completion_input_ids.shape + device = completion_input_ids.device + + num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt + + indices = torch.arange(completion_len, device=device).unsqueeze(0) + shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) + + non_padding_mask = (completion_input_ids != pad_token_id) + + final_mask = shift_mask & non_padding_mask + + return final_mask + +def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: + """Move all padding tokens in each sequence to the right.""" + mask = (tensor != pad_id) + # stable=True since the binary mask is unordered. + sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) + packed_tensor = torch.gather(tensor, 1, sorted_indices) + return packed_tensor + +def align_logprobs_with_mask( + logprob_tensor: torch.Tensor, + attention_mask: torch.Tensor, + pad_value: float = 0.0 +) -> torch.Tensor: + """Align a log probability tensor with a given attention mask.""" + + device = logprob_tensor.device + batch_size, logprob_seq_len = logprob_tensor.shape + mask_seq_len = attention_mask.shape[1] + + padded_logprobs = torch.full( + attention_mask.shape, + fill_value=pad_value, + dtype=logprob_tensor.dtype, + device=device + ) + + left_pad_counts = torch.argmax(attention_mask, dim=1) + + cols = torch.arange(logprob_seq_len, device=device) + dest_indices = left_pad_counts.unsqueeze(1) + cols + + # Destination row indices, shape [batch_size, logprob_seq_len]. + row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) + + # Keep only in-bounds destinations, then scatter via advanced indexing. + valid_mask = dest_indices < mask_seq_len + valid_rows = row_indices[valid_mask] + valid_cols = dest_indices[valid_mask] + valid_vals = logprob_tensor[valid_mask] + padded_logprobs[valid_rows, valid_cols] = valid_vals + + return padded_logprobs + +def align_completion_tool_mask( + tool_mask: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" + if tool_mask is None: + return completion_mask + if tool_mask.shape[0] != completion_mask.shape[0]: + raise ValueError("tool_mask batch size must match completion_mask batch size.") + + tool_mask = tool_mask.to(device=completion_mask.device) + if tool_mask.shape == completion_mask.shape: + aligned_tool_mask = tool_mask + else: + aligned_tool_mask = align_logprobs_with_mask( + tool_mask, + completion_mask, + pad_value=0, + ) + return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) + +def autotune_batch_and_chunks( + total_input_rows, + seq_len, + hidden_size, + vocab_size, + dtype_bytes=16, + multiplier=None +): + if multiplier is None: + final_m = max(4, seq_len // 4096) + else: + final_m = multiplier + + if torch.cuda.is_available(): + free_bytes, _ = torch.cuda.mem_get_info() + limit_gb = (free_bytes / (1024**3))*.80 + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + # XPU: estimate free memory as total - reserved. + total_mem = torch.xpu.get_device_properties(0).total_memory + reserved_mem = torch.xpu.memory_reserved() + free_bytes = total_mem - reserved_mem + limit_gb = (free_bytes / (1024**3)) * 0.80 + else: + # Fallback: assume 8GB available. + limit_gb = 8.0 + + bytes_to_gb = 1024**3 + + b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) + + hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb + + base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb + logits_gb = base_logits / final_m + + total_mem_gb = hidden_gb + logits_gb + + valid_mask = total_mem_gb <= limit_gb + valid_indices = torch.nonzero(valid_mask, as_tuple=False) + + if valid_indices.shape[0] == 0: + #This means your GPU will OOM + return 4, final_m + + best_idx = valid_indices[0].item() + final_b = int(b_vals[best_idx].item()) + + return final_b, final_m + +def sanitize_logprob(logprob): + """Local port of trl.scripts.vllm_serve.sanitize_logprob. + Filters NaN logprobs from vLLM outputs.""" + value = logprob.logprob + if math.isnan(value): + logging.getLogger(__name__).warning( + f"Generated NaN logprob, token logprob '{logprob}' will be ignored" + ) + return None + return value +@dataclass +class UnslothSFTConfig(SFTConfig): + """ + + Configuration class for the [`SFTTrainer`]. + + This class includes only the parameters that are specific to SFT training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`SFTTrainer`] is provided as a string. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to + [`~transformers.AutoModelForCausalLM.from_pretrained`] and + [`~transformers.AutoProcessor.from_pretrained`]. + router_aux_loss_coef (`float`, *optional*, defaults to `0.001`): + Coefficient of the load-balancing auxiliary loss. Only has an effect when training a Mixture-of-Experts + (MoE) model; for other models it does nothing. The auxiliary loss is added to the training loss with this + weight. Set to `0.0` to disable it. + chat_template_path (`str`, *optional*): + If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory + or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must + ensure that any special tokens referenced in the template are added to the tokenizer and that the model's + embedding layer is resized accordingly. + + > Parameters that control the data preprocessing + + dataset_text_field (`str`, *optional*, defaults to `"text"`): + Name of the column that contains text data in the dataset. + dataset_kwargs (`dict[str, Any]`, *optional*): + Dictionary of optional keyword arguments for the dataset preparation. The only supported key is + `skip_prepare_dataset`. When the model is a VLM, `skip_prepare_dataset` is automatically treated as `True` + regardless of the provided value, since preprocessing is done on the fly. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + eos_token (`str`, *optional*): + Token used to indicate the end of a turn or sequence. If `None`, it defaults to + `processing_class.eos_token`. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the left + or right depending on `truncation_mode`. If `None`, no truncation is applied. When packing is enabled, + this value sets the sequence length. + truncation_mode (`str`, *optional*, defaults to `"keep_start"`): + Truncation mode to use when the sequence exceeds `max_length`. The only supported value is + `"keep_start"`. The `"keep_end"` value is deprecated and will be removed in v2.0.0. + shuffle_dataset (`bool`, *optional*, defaults to `False`): + Whether to shuffle the dataset. + packing (`bool`, *optional*, defaults to `False`): + Whether to group multiple sequences into fixed-length blocks to improve computational efficiency and reduce + padding. Uses `max_length` to define sequence length. + packing_strategy (`str`, *optional*, defaults to `"bfd"`): + Strategy for packing sequences. Can be `"bfd"` (best-fit decreasing, truncates overflow), `"bfd_split"` + (best-fit decreasing, splits overflow sequences), or `"wrapped"` (aggressive, cuts mid-sequence). + padding_free (`bool`, *optional*, defaults to `False`): + Whether to perform forward passes without padding by flattening all sequences in the batch into a single + continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only + supported with the FlashAttention 2 or 3, which can efficiently handle the flattened batch structure. When + packing is enabled with strategy `"bfd"`, padding-free is enabled, regardless of the value of this + parameter. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + eval_packing (`bool`, *optional*): + Whether to pack the eval dataset. If `None`, uses the same value as `packing`. + + > Parameters that control the training + + completion_only_loss (`bool`, *optional*): + Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is computed + only on the completion, which is supported only for [prompt-completion](#prompt-completion) datasets. If + `False`, loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: + loss is computed on the completion for [prompt-completion](#prompt-completion) datasets, and on the full + sequence for [language modeling](#language-modeling) datasets. + assistant_only_loss (`bool`, *optional*, defaults to `False`): + Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is computed only + on the assistant responses, which is supported only for [conversational](#conversational) datasets. If + `False`, loss is computed on the entire sequence. + loss_type (`str`, *optional*, defaults to `"chunked_nll"`): + Type of loss to use. When left unset, it defaults to `"chunked_nll"`, except when `use_liger_kernel=True`, + in which case it defaults to `"nll"`. Possible values are: + + - `"nll"`: standard negative log-likelihood. + - `"dft"`: Dynamic Fine-Tuning, as described in + [this paper](https://huggingface.co/papers/2508.05629). + - `"chunked_nll"`: same math as `"nll"`, but the `lm_head` projection is computed on non-ignored tokens + only (positions with `labels == -100` are dropped before the matmul) and the cross-entropy is processed + in chunks of tokens to reduce peak activation memory. Not compatible with `use_liger_kernel`. + + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + + > Deprecated parameters + + pad_token: + + + + Parameter `pad_token` is deprecated and will be removed in version v2.0.0. Set `tokenizer.pad_token` + directly and pass it as `processing_class` to the trainer instead. + + + + > [!NOTE] + > These parameters have default values different from [`~transformers.TrainingArguments`]: + > - `logging_steps`: Defaults to `10` instead of `500`. + > - `gradient_checkpointing`: Defaults to `True` instead of `False`. + > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. + > - `learning_rate`: Defaults to `2e-5` instead of `5e-5`. + + """ + vllm_sampling_params: Optional[Any] = field( + default = None, + metadata = {'help': 'vLLM SamplingParams'}, + ) + unsloth_num_chunks : Optional[int] = field( + default = -1, + metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, + ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {'help': 'Multiplier for chunked logit computations.'}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, + ) + max_seq_length : Optional[int] = field( + default = None, + metadata = {'help': 'Maximum sequence length to truncate to.'}, + ) + def __init__( + self, + output_dir = None, + per_device_train_batch_size = 4, + num_train_epochs = 3.0, + max_steps = -1, + learning_rate = 5e-05, + lr_scheduler_type = 'linear', + lr_scheduler_kwargs = None, + warmup_steps = 0.1, + optim = 'adamw_8bit', + optim_args = None, + weight_decay = 0.001, + adam_beta1 = 0.9, + adam_beta2 = 0.999, + adam_epsilon = 1e-08, + optim_target_modules = None, + gradient_accumulation_steps = 2, + average_tokens_across_devices = True, + max_grad_norm = 1.0, + label_smoothing_factor = 0.0, + bf16 = False, + fp16 = False, + bf16_full_eval = False, + fp16_full_eval = False, + tf32 = None, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = None, + torch_compile = False, + torch_compile_backend = None, + torch_compile_mode = None, + use_liger_kernel = False, + liger_kernel_config = None, + use_cache = False, + neftune_noise_alpha = None, + torch_empty_cache_steps = 250, + auto_find_batch_size = False, + logging_strategy = 'steps', + logging_steps = 1, + logging_first_step = False, + log_on_each_node = True, + logging_nan_inf_filter = False, + include_num_input_tokens_seen = False, + log_level = 'passive', + log_level_replica = 'warning', + disable_tqdm = None, + report_to = 'none', + run_name = None, + project = 'huggingface', + trackio_space_id = 'trackio', + eval_strategy = 'no', + eval_steps = None, + eval_delay = 0, + per_device_eval_batch_size = 4, + prediction_loss_only = False, + eval_on_start = False, + eval_do_concat_batches = True, + eval_use_gather_object = False, + eval_accumulation_steps = 2, + batch_eval_metrics = False, + save_only_model = False, + save_strategy = 'steps', + save_steps = 500, + save_on_each_node = False, + save_total_limit = None, + enable_jit_checkpoint = False, + push_to_hub = False, + hub_token = None, + hub_private_repo = None, + hub_model_id = None, + hub_strategy = 'every_save', + hub_always_push = False, + hub_revision = None, + load_best_model_at_end = False, + metric_for_best_model = None, + greater_is_better = None, + ignore_data_skip = False, + restore_callback_states_from_checkpoint = False, + full_determinism = False, + seed = 3407, + data_seed = 3407, + use_cpu = False, + accelerator_config = None, + parallelism_config = None, + dataloader_drop_last = False, + dataloader_num_workers = 0, + dataloader_pin_memory = True, + dataloader_persistent_workers = False, + dataloader_prefetch_factor = None, + remove_unused_columns = True, + label_names = None, + train_sampling_strategy = 'random', + length_column_name = 'length', + ddp_find_unused_parameters = None, + ddp_bucket_cap_mb = None, + ddp_broadcast_buffers = None, + ddp_backend = None, + ddp_timeout = 1800, + fsdp = None, + fsdp_config = None, + deepspeed = None, + debug = '', + skip_memory_metrics = True, + do_train = False, + do_eval = False, + do_predict = False, + resume_from_checkpoint = None, + warmup_ratio = None, + logging_dir = None, + local_rank = -1, + model_init_kwargs = None, + router_aux_loss_coef = 0.001, + trust_remote_code = False, + chat_template_path = None, + dataset_text_field = 'text', + dataset_kwargs = None, + dataset_num_proc = None, + eos_token = None, + max_length = 1024, + truncation_mode = 'keep_start', + shuffle_dataset = False, + packing = False, + packing_strategy = 'bfd', + padding_free = None, + pad_to_multiple_of = None, + eval_packing = None, + completion_only_loss = None, + assistant_only_loss = False, + loss_type = None, + activation_offloading = False, + pad_token = None, + vllm_sampling_params = None, + unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, + max_seq_length = None, + **kwargs, + ): + if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') + if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') + if num_train_epochs is None: + num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override + if output_dir is None and save_strategy == 'steps' and save_steps == 500: + output_dir = 'unsloth_training_checkpoints' + save_strategy = 'no' + import multiprocessing as _mp + if dataset_num_proc is None: + if _mp.get_start_method() != 'fork': + dataset_num_proc = None + else: + import psutil + dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64) + memory_gb_left = psutil.virtual_memory().available / (1024**3) + if memory_gb_left <= 2: dataset_num_proc = 1 + else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left)) + if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': + from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION + if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: + from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE + pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE + + + super().__init__( + output_dir = output_dir, + per_device_train_batch_size = per_device_train_batch_size, + num_train_epochs = num_train_epochs, + max_steps = max_steps, + learning_rate = learning_rate, + lr_scheduler_type = lr_scheduler_type, + lr_scheduler_kwargs = lr_scheduler_kwargs, + warmup_steps = warmup_steps, + optim = optim, + optim_args = optim_args, + weight_decay = weight_decay, + adam_beta1 = adam_beta1, + adam_beta2 = adam_beta2, + adam_epsilon = adam_epsilon, + optim_target_modules = optim_target_modules, + gradient_accumulation_steps = gradient_accumulation_steps, + average_tokens_across_devices = average_tokens_across_devices, + max_grad_norm = max_grad_norm, + label_smoothing_factor = label_smoothing_factor, + bf16 = bf16, + fp16 = fp16, + bf16_full_eval = bf16_full_eval, + fp16_full_eval = fp16_full_eval, + tf32 = tf32, + gradient_checkpointing = gradient_checkpointing, + gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, + torch_compile = torch_compile, + torch_compile_backend = torch_compile_backend, + torch_compile_mode = torch_compile_mode, + use_liger_kernel = use_liger_kernel, + liger_kernel_config = liger_kernel_config, + use_cache = use_cache, + neftune_noise_alpha = neftune_noise_alpha, + torch_empty_cache_steps = torch_empty_cache_steps, + auto_find_batch_size = auto_find_batch_size, + logging_strategy = logging_strategy, + logging_steps = logging_steps, + logging_first_step = logging_first_step, + log_on_each_node = log_on_each_node, + logging_nan_inf_filter = logging_nan_inf_filter, + include_num_input_tokens_seen = include_num_input_tokens_seen, + log_level = log_level, + log_level_replica = log_level_replica, + disable_tqdm = disable_tqdm, + report_to = report_to, + run_name = run_name, + project = project, + trackio_space_id = trackio_space_id, + eval_strategy = eval_strategy, + eval_steps = eval_steps, + eval_delay = eval_delay, + per_device_eval_batch_size = per_device_eval_batch_size, + prediction_loss_only = prediction_loss_only, + eval_on_start = eval_on_start, + eval_do_concat_batches = eval_do_concat_batches, + eval_use_gather_object = eval_use_gather_object, + eval_accumulation_steps = eval_accumulation_steps, + batch_eval_metrics = batch_eval_metrics, + save_only_model = save_only_model, + save_strategy = save_strategy, + save_steps = save_steps, + save_on_each_node = save_on_each_node, + save_total_limit = save_total_limit, + enable_jit_checkpoint = enable_jit_checkpoint, + push_to_hub = push_to_hub, + hub_token = hub_token, + hub_private_repo = hub_private_repo, + hub_model_id = hub_model_id, + hub_strategy = hub_strategy, + hub_always_push = hub_always_push, + hub_revision = hub_revision, + load_best_model_at_end = load_best_model_at_end, + metric_for_best_model = metric_for_best_model, + greater_is_better = greater_is_better, + ignore_data_skip = ignore_data_skip, + restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, + full_determinism = full_determinism, + seed = seed, + data_seed = data_seed, + use_cpu = use_cpu, + accelerator_config = accelerator_config, + parallelism_config = parallelism_config, + dataloader_drop_last = dataloader_drop_last, + dataloader_num_workers = dataloader_num_workers, + dataloader_pin_memory = dataloader_pin_memory, + dataloader_persistent_workers = dataloader_persistent_workers, + dataloader_prefetch_factor = dataloader_prefetch_factor, + remove_unused_columns = remove_unused_columns, + label_names = label_names, + train_sampling_strategy = train_sampling_strategy, + length_column_name = length_column_name, + ddp_find_unused_parameters = ddp_find_unused_parameters, + ddp_bucket_cap_mb = ddp_bucket_cap_mb, + ddp_broadcast_buffers = ddp_broadcast_buffers, + ddp_backend = ddp_backend, + ddp_timeout = ddp_timeout, + fsdp = fsdp, + fsdp_config = fsdp_config, + deepspeed = deepspeed, + debug = debug, + skip_memory_metrics = skip_memory_metrics, + do_train = do_train, + do_eval = do_eval, + do_predict = do_predict, + resume_from_checkpoint = resume_from_checkpoint, + warmup_ratio = warmup_ratio, + logging_dir = logging_dir, + local_rank = local_rank, + model_init_kwargs = model_init_kwargs, + router_aux_loss_coef = router_aux_loss_coef, + trust_remote_code = trust_remote_code, + chat_template_path = chat_template_path, + dataset_text_field = dataset_text_field, + dataset_kwargs = dataset_kwargs, + dataset_num_proc = dataset_num_proc, + eos_token = eos_token, + max_length = max_length, + truncation_mode = truncation_mode, + shuffle_dataset = shuffle_dataset, + packing = packing, + packing_strategy = packing_strategy, + padding_free = padding_free, + pad_to_multiple_of = pad_to_multiple_of, + eval_packing = eval_packing, + completion_only_loss = completion_only_loss, + assistant_only_loss = assistant_only_loss, + loss_type = loss_type, + activation_offloading = activation_offloading, + pad_token = pad_token,**kwargs) + self.vllm_sampling_params = vllm_sampling_params + self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier + self.max_seq_length = max_seq_length + # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+ + if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: + if 'use_reentrant' in self.gradient_checkpointing_kwargs: + del self.gradient_checkpointing_kwargs['use_reentrant'] + +pass + +class _UnslothSFTTrainer(_BaseTrainer): + """""" + + _tag_names = ["trl", "sft"] + _name = "SFT" + + def __init__( + self, + model: "str | PreTrainedModel | PeftModel", + args: SFTConfig | TrainingArguments | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + compute_loss_func: Callable | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: "PeftConfig | None" = None, + formatting_func: Callable[[dict], str] | None = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = SFTConfig(f"{model_name}-SFT") + elif isinstance(args, TrainingArguments) and not isinstance(args, SFTConfig): + dict_args = args.to_dict() + dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token + if Version(transformers.__version__) < Version("5.0.0"): + dict_args.pop("push_to_hub_token", None) + args = SFTConfig(**dict_args) + + if train_dataset is None: + raise ValueError("`train_dataset` is required") + elif isinstance(train_dataset, IterableDataset): + # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate + # batches from multiple processes, leading to mismatch errors. + if args.accelerator_config.dispatch_batches is True: + logger.warning( + "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " + "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " + "`dispatch_batches` in `SFTConfig` or set it to `False`." + ) + args.accelerator_config.dispatch_batches = False + + # Model + if isinstance(model, str): + model_init_kwargs = args.model_init_kwargs or {} + # Distributed training requires device_map=None ["auto" fails] + if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: + model_init_kwargs["device_map"] = None + model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) + model = create_model_from_path(model, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `SFTConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do + _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained( + get_config_model_id(model.config), trust_remote_code=args.trust_remote_code + ) + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + self._tokenizer = processing_class.tokenizer + self._is_vlm = True + elif isinstance(processing_class, PreTrainedTokenizerBase): + self._tokenizer = processing_class + self._is_vlm = False + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if args.eos_token is not None: + if args.eos_token not in self._tokenizer.get_vocab(): + raise ValueError( + f"The specified `eos_token` ('{args.eos_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " + "in the vocabulary before using it as an EOS token." + ) + self._tokenizer.eos_token = args.eos_token + + if args.chat_template_path is not None: + if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): + with open(args.chat_template_path, encoding="utf-8") as chat_template_file: + processing_class.chat_template = chat_template_file.read() + added_tokens = [] + else: + model, processing_class, added_tokens = clone_chat_template( + model, processing_class, args.chat_template_path + ) + else: + added_tokens = [] + + # Catch some wrong configurations related to VLMs + if self._is_vlm and args.packing: + raise ValueError( + "Packing is not supported for vision-language models. Please set `packing=False` in the SFTConfig." + ) + if self._is_vlm and args.padding_free: + raise ValueError( + "Padding-free training is yet not supported for vision-language models. Please set " + "`padding_free=False` in the `SFTConfig`." + ) + if self._is_vlm and args.assistant_only_loss: + raise ValueError( + "Assistant-only loss is not yet supported for vision-language models. Please set " + "`assistant_only_loss=False` in the `SFTConfig`." + ) + if self._is_vlm and args.max_length is not None and args.truncation_mode == "keep_end": + raise ValueError( + "truncation_mode='keep_end' is not supported for vision-language models. Image tokens reside " + "inside the prompt portion of the sequence; depending on the example, keep_end may silently " + "drop them, causing pixel_values to be forwarded to the model with no corresponding visual " + "tokens in input_ids. Use truncation_mode='keep_start' (the default) or set max_length=None." + ) + + # PEFT + if False: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if is_peft_model(model): + raise ValueError( + "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " + "and unload the existing adapter, save the resulting base model, and then pass that base model along " + "with the new `peft_config` to the trainer." + ) + if added_tokens: + # Ensure that the added tokens are trainable + if peft_config.trainable_token_indices is None: + peft_config.trainable_token_indices = {"embed_tokens": added_tokens} + elif "embed_tokens" not in peft_config.trainable_token_indices: + peft_config.trainable_token_indices["embed_tokens"] = added_tokens + else: + peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) + # Ensure that the lm_head is trainable + if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: + logger.warning( + "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " + "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " + "tokens, leading to degraded generation quality. To fix this, add " + "`modules_to_save=['lm_head']` to your PEFT configuration." + ) + + if peft_config.modules_to_save is None: + peft_config.modules_to_save = ["lm_head"] + else: + peft_config.modules_to_save.append("lm_head") + # Create PEFT model + # ZeRO-3 + PEFT for non-quantized models: + # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. + # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, + # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. + # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. + # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for + # non-quantized bf16 training. + # - See: + # - TRL issue: https://github.com/huggingface/trl/issues/6089 + # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 + # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg + get_peft_model_kwargs = {} + if ( + args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and not _is_quantized_model + and Version(peft.__version__) >= Version("0.12.0") + ): + get_peft_model_kwargs["autocast_adapter_dtype"] = False + model = get_peft_model(model, peft_config, **get_peft_model_kwargs) + + # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see + # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703 + if ( + is_peft_model(model) + and args.deepspeed_plugin is not None + and args.deepspeed_plugin.zero_stage == 3 + and args.gradient_checkpointing + ): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") + if use_reentrant is False: + logger.warning( + "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " + "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " + "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." + ) + args.gradient_checkpointing_kwargs["use_reentrant"] = True + + # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally + # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 + if is_peft_model(model) and args.gradient_checkpointing: + model.enable_input_require_grads() + + # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the + # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by + # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for + # quantized models. See: https://github.com/huggingface/peft/issues/2889 + if _is_quantized_model: + for param in model.parameters(): + if param.requires_grad: + param.data = param.data.to(torch.bfloat16) + + # In Prompt Tuning a small set of trainable virtual tokens [continuous prompt embeddings] is prepended to the + # input. We store the number of these tokens so we can account for them correctly when calculating accuracy. + self.num_virtual_tokens = 0 + if is_peft_model(model): + if model.active_adapter in model.peft_config: + peft_model_config = model.peft_config[model.active_adapter] + self.num_virtual_tokens = getattr(peft_model_config, "num_virtual_tokens", 0) + + # Data collator + # BFD packing requires padding-free mode; otherwise, the collator outputs padded attention masks, causing + # FlashAttention to ignore position_ids and recompute them incorrectly from the padded attention mask. + self.padding_free = args.padding_free or (args.packing and args.packing_strategy in {"bfd", "bfd_split"}) + use_flash_attention = model.config._attn_implementation in FLASH_ATTENTION_VARIANTS + if self.padding_free: + if data_collator is not None: + raise ValueError("Passing a custom data collator is not supported when using padding-free.") + if args.packing and args.packing_strategy == "wrapped": + logger.warning( + "You are passing `padding_free=True` with the 'wrapped' packing strategy, which is not " + "recommended. Please refer to the documentation to understand why this is not recommended." + ) + if not use_flash_attention: + logger.warning( + "Padding-free training is enabled, but the attention implementation is not set to a supported " + "Flash Attention variant. Padding-free training flattens batches into a single sequence, and only " + "the following implementations are known to reliably support this: " + f"{', '.join(sorted(FLASH_ATTENTION_VARIANTS))}. Using other implementations may lead to " + "unexpected behavior. To ensure compatibility, set `attn_implementation` in the model " + "configuration to one of these supported options or verify that your attention mechanism can " + "handle flattened sequences." + ) + # Decide whether to use completion-only loss: if not specified, then it is set to True if the dataset format + # is prompt-completion, and False if the dataset format is language modeling. + dataset_sample = next(iter(train_dataset)) + if args.completion_only_loss is None: + self.completion_only_loss = "prompt" in dataset_sample and "completion" in dataset_sample + else: + self.completion_only_loss = args.completion_only_loss + + self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample + # Unsloth: override _is_vlm for VLM models that pass a bare tokenizer + if not self._is_vlm and self._is_vision_dataset: + _m = model + if hasattr(_m, "model"): _m = _m.model + if hasattr(getattr(_m, "config", None), "vision_config") or \ + _m.__class__.__name__.endswith("ForConditionalGeneration"): + self._is_vlm = True + if self._is_vision_dataset and not self._is_vlm: + raise ValueError( + "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " + "model does not seem to be a vision-language model. Please check your model and dataset." + ) + + if data_collator is None and not self._is_vision_dataset: + # Get the pad token: if not provided, use the one from the processing class or the eos token + # if the processing class does not have a pad token. + pad_token = args.pad_token or self._tokenizer.pad_token or self._tokenizer.eos_token + if pad_token not in self._tokenizer.get_vocab(): + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + self._tokenizer.pad_token = pad_token + data_collator = DataCollatorForLanguageModeling( + pad_token_id=self._tokenizer.pad_token_id, + max_length=None if self.padding_free else args.max_length, + truncation_mode=args.truncation_mode, + padding_free=self.padding_free, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + elif data_collator is None and self._is_vision_dataset: + data_collator = DataCollatorForVisionLanguageModeling( + processor=processing_class, + max_length=args.max_length, + completion_only_loss=self.completion_only_loss, + pad_to_multiple_of=args.pad_to_multiple_of, + dataset_text_field=args.dataset_text_field, + ) + + if args.packing and args.packing_strategy in {"bfd", "bfd_split"} and not use_flash_attention: + logger.warning( + "You are using packing, but the attention implementation is not set to a supported Flash Attention " + "variant. Packing gathers multiple samples into a single sequence, and only the following " + f"implementations are known to reliably support this: {', '.join(sorted(FLASH_ATTENTION_VARIANTS))}. " + "Using other implementations may lead to cross-contamination between samples. To avoid this, either " + "disable packing by setting `packing=False`, or set `attn_implementation` in the model configuration " + "to one of these supported options." + ) + if args.assistant_only_loss and not is_conversational(dataset_sample): + raise ValueError( + "You set `assistant_only_loss=True`, but the dataset is not conversational. This option is only " + "supported for conversational datasets." + ) + + # When assistant_only_loss is enabled, swap in a training chat template with {% generation %} markers + # if the current template doesn't already have them. + if args.assistant_only_loss and not has_generation_markers(processing_class.chat_template): + self.chat_template = get_training_chat_template(processing_class) + else: + self.chat_template = None + + # A template can define generation markers and still attribute the assistant's end-of-turn token to the next + # message, leaving it out of the assistant mask so the model is never trained to stop. + if args.assistant_only_loss and not is_chat_template_stop_token_trained( + processing_class, chat_template=self.chat_template + ): + logger.warning( + "The chat template does not include the assistant turn's end-of-turn token in the loss mask; " + "the model may not learn to stop." + ) + + # Dataset + if self.padding_free and not args.packing and args.max_length is not None and not self._is_vision_dataset: + raise ValueError( + "When `padding_free=True` without packing, `max_length` is not enforced. Either enable packing " + "(e.g., `packing=True, packing_strategy='bfd'`), provide already truncated inputs, or set " + "`max_length=None`." + ) + # Skip dataset preparation if `skip_prepare_dataset=True` in `dataset_kwargs`, or if it's a VLM, where + # preprocessing [e.g., image-to-pixel conversion] is too costly and done on the fly instead. + self._skip_prepare_dataset = ( + args.dataset_kwargs is not None + and args.dataset_kwargs.get("skip_prepare_dataset", False) + or self._is_vision_dataset + ) + # Kept on the instance so that `evaluate` can preprocess freshly-passed eval datasets the same way. + self._formatting_func = formatting_func + eval_datasets = ( + eval_dataset if isinstance(eval_dataset, dict) else {"eval": eval_dataset} if eval_dataset else {} + ) + self._reject_skip_prepare_without_labels({"train": train_dataset, **eval_datasets}, data_collator) + if not self._skip_prepare_dataset: + if self.completion_only_loss and formatting_func: + raise ValueError( + "A formatting function was provided while `completion_only_loss=True`, which is incompatible. " + "Using a formatter converts the dataset to a language modeling type, conflicting with " + "completion-only loss. To resolve this, apply your formatting function before passing the " + "dataset, or disable `completion_only_loss` in `SFTConfig`." + ) + self._unsloth_model_ref = model + train_dataset = self._prepare_dataset( + train_dataset, processing_class, args, args.packing, formatting_func, "train" + ) + if eval_dataset is not None: + packing = args.packing if args.eval_packing is None else args.eval_packing + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, packing, formatting_func, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset( + eval_dataset, processing_class, args, packing, formatting_func, "eval" + ) + + # Loss function + if not args.use_liger_kernel: # liger supports dft loss by just passing use_token_scaling=True + if args.loss_type == "nll": + pass # use the default loss + elif args.loss_type == "dft": + if compute_loss_func is not None: + raise ValueError( + "You passed a `compute_loss_func` together with `loss_type='dft'` to the `SFTTrainer`. " + "When using `loss_type='dft'`, the loss function is internally set to the DFT loss, so " + "passing a `compute_loss_func` is not allowed." + ) + compute_loss_func = dft_loss + elif args.loss_type == "chunked_nll": + # Same math as `"nll"` but the `lm_head` matmul is skipped on ignored tokens and the CE is computed in + # chunks of tokens. Implemented by patching the model's forward before `super[].__init__` so accelerate + # wraps the patched forward. + # For PEFT, patch the inner causal LM rather than the `PeftModel` wrapper. LoRA / IA³ / + # `modules_to_save` adapters live in the module tree, so they're hit even when we bypass + # `PeftModel.forward`. Prompt-learning variants need `PeftModel.forward` to run first [to inject + # virtual tokens], then it delegates into the patched inner forward. + target = model.get_base_model() if is_peft_model(model) else model + # The chunked path reads the output projection weight directly, which would silently drop the + # adapter delta [and starve its parameters of gradients] if the head itself is a PEFT tuner layer. + if is_peft_model(model): + from peft.tuners.tuners_utils import BaseTunerLayer + + if isinstance(target.get_output_embeddings(), BaseTunerLayer): + raise ValueError( + "`loss_type='chunked_nll'` is not supported when `lm_head` is wrapped by a PEFT adapter " + "(e.g. `target_modules='all-linear'` or explicitly including `'lm_head'`). Either remove " + "`lm_head` from `target_modules`, or switch to `loss_type='nll'`. If this is a real use " + "case for you, please open an issue at https://github.com/huggingface/trl/issues." + ) + _patch_chunked_ce_lm_head(target, chunk_size=_CHUNKED_LM_HEAD_CHUNK_SIZE, is_vlm=self._is_vlm) + else: + raise ValueError( + f"Invalid `loss_type` {args.loss_type} passed. Supported values are 'nll', 'dft', and " + "'chunked_nll'." + ) + elif args.loss_type == "chunked_nll": + raise ValueError("`loss_type='chunked_nll'` is not compatible with `use_liger_kernel=True`.") + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): + args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_loss_func=compute_loss_func, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # MoE load-balancing auxiliary loss, applied to Mixture-of-Experts models [no effect otherwise] + text_config = model.config.get_text_config() + is_moe = getattr(text_config, "output_router_logits", None) is not None + self.aux_loss_enabled = is_moe and self.args.router_aux_loss_coef != 0.0 + if is_moe: + # The native and chunked forwards add the aux loss from the model config, so keep the config in sync with + # the coef: enable it [and propagate the coef] when non-zero, disable it otherwise. This overrides any + # `output_router_logits` the model was loaded with, so `router_aux_loss_coef=0.0` reliably turns it off. + text_config.output_router_logits = self.aux_loss_enabled + text_config.router_aux_loss_coef = self.args.router_aux_loss_coef + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + def _tokenize( + self, + processing_class: PreTrainedTokenizerBase | ProcessorMixin, + input: str | list, + **kwargs, + ) -> dict[str, list]: + """Tokenize a single example for dataset preprocessing. + + Dispatches to `apply_chat_template` for conversational input (list of message dicts) and to `__call__` for + non-conversational input (str). For VLMs, normalizes the batch dimension that processors emit even for single + examples. + + Args: + processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`]): + The tokenizer or processor to use. + input (`str` or `list`): + A string for non-conversational input, or a list of message dicts for conversational input. + **kwargs: + Forwarded to `apply_chat_template` (e.g. `add_generation_prompt`, `return_assistant_tokens_mask`). + + Returns: + `dict` with at least an `"input_ids"` key mapping to a flat `list[int]`. + """ + if isinstance(input, list): # conversational: list of message dicts + if self._is_vlm: + input = prepare_multimodal_messages(input) + result = processing_class.apply_chat_template( + input, tokenize=True, return_dict=True, chat_template=self.chat_template, **kwargs + ) + else: # non-conversational: plain text string + result = processing_class(text=input) + # VLMs emit a batch dimension even for single examples; unwrap it + if self._is_vlm: + return {k: v[0] for k, v in result.items()} + return result + + def _prepare_dataset( + self, + dataset: Union[Dataset, IterableDataset], + processing_class, + args, + packing: bool, + formatting_func: Optional[Callable[[dict], str]], + dataset_name: str, + ) -> Union[Dataset, IterableDataset]: + # All Unsloth Zoo code licensed under LGPLv3 + try: + if isinstance(dataset, ConstantLengthDataset): return dataset + except: + pass + + map_kwargs = {} + use_desc = isinstance(dataset, Dataset) + is_vlm = hasattr(processing_class, "tokenizer") + tokenizer = processing_class + if is_vlm: tokenizer = processing_class.tokenizer + + # Detect whether the model's module needs token_type_ids when training + import sys as _sys + _needs_token_type_ids = False + # Split to avoid compiler substring match on masking_utils names + _ccm = 'create_' + 'causal_mask_mapping' + _model = getattr(self, '_unsloth_model_ref', None) or getattr(self, 'model', None) + if _model is not None: + for _m in (_model, getattr(_model, 'model', None)): + if _m is None: continue + _mod = _sys.modules.get(type(_m).__module__) + if _mod is not None and hasattr(_mod, _ccm): + _needs_token_type_ids = True + break + + if not _needs_token_type_ids: + # Fallback: model not yet available, check processor class MRO + for _base in type(processing_class).__mro__: + _base_mod = getattr(_base, '__module__', '') + if 'transformers.models.' in _base_mod: + _modeling_mod = _base_mod.replace('.processing_', '.modeling_') + _mod = _sys.modules.get(_modeling_mod) + if _mod is not None and hasattr(_mod, _ccm): + _needs_token_type_ids = True + break + if _needs_token_type_ids and hasattr(args, 'remove_unused_columns'): + args.remove_unused_columns = False + + # Get max length + max_seq_length = getattr(args, "max_length", 0) + if max_seq_length == 0: max_seq_length = getattr(args, "max_seq_length", 0) + if max_seq_length == 0: max_seq_length = getattr(self, "max_seq_length", 0) + if max_seq_length == 0: max_seq_length = getattr(self, "max_seq", 0) + if max_seq_length == 0: raise RuntimeError("Unsloth: max_seq_length is 0! Please specify one!") + dataset_text_field = getattr(args, "dataset_text_field", "text") + do_truncation = max_seq_length != 0 + do_formatting_func = False + do_tokenize = True + do_prompt_completion = False + + # Get correct column names + column_names = set(next(iter(dataset)).keys()) + used_column_names = ["input_ids"] + if "attention_mask" in column_names: + used_column_names.append("attention_mask") + if _needs_token_type_ids: + used_column_names.append("token_type_ids") + + # Skip tokenization if already tokenized; just set the data collator + from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling + if "labels" in column_names: + # Most likely forgot data collator! + if is_vlm and not hasattr(tokenizer, "pad"): + raise RuntimeError(f"Unsloth: {processing_class.__class__} does not have .pad!") + self.data_collator = DataCollatorForSeq2Seq(tokenizer) + used_column_names.append("labels") + do_tokenize = False + elif "input_ids" in column_names: + if is_vlm and not hasattr(tokenizer, "pad"): + raise RuntimeError(f"Unsloth: {processing_class.__class__} does not have .pad!") + self.data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) + do_tokenize = False + elif "prompt" in column_names and "completion" in column_names: + # Prompt/completion dataset (used with completion_only_loss). + # TRL's __init__ already set self.data_collator for completion_only_loss + # before calling us -- we must NOT overwrite it here. + do_prompt_completion = True + used_column_names.append("completion_mask") + elif dataset_text_field not in column_names: + do_formatting_func = True + if formatting_func is None: + raise RuntimeError("Unsloth: You must specify a `formatting_func`") + pass + + if do_tokenize: + if do_formatting_func: + test_text = formatting_func(next(iter(dataset))) + if not isinstance(test_text, list): + raise ValueError( + "Unsloth: The `formatting_func` should return a list of processed strings." + ) + test_text = test_text[0] + elif do_prompt_completion: + _first_ex = next(iter(dataset)) + try: + from trl import is_conversational as _sft_is_conversational + except ImportError: + def _sft_is_conversational(example): + for key in ("prompt", "completion", "messages"): + val = example.get(key) + if isinstance(val, list) and val and isinstance(val[0], dict): + if "role" in val[0] and "content" in val[0]: + return True + return False + _is_conv = _sft_is_conversational(_first_ex) + if not _is_conv: + test_text = _first_ex["prompt"] + else: + test_text = None # chat template handles BOS + else: + test_text = next(iter(dataset))[dataset_text_field][0] + + chat_template = getattr(processing_class, 'chat_template', '') + if chat_template == '' and is_vlm: + chat_template = getattr(tokenizer, 'chat_template', '') + if chat_template is None: + chat_template = '' + + # Detect double BOS so we can drop the duplicate + add_special_tokens = True + bos_token_1 = getattr(processing_class, 'bos_token', None) + bos_token_2 = getattr(tokenizer, 'bos_token', None) + bos_token = bos_token_1 or bos_token_2 + + if bos_token is not None: + if (test_text is not None and test_text.startswith(bos_token)) or bos_token in chat_template: + add_special_tokens = False + print("Unsloth: We found double BOS tokens - we shall remove one automatically.") + pass + + def _tokenize(example): + return tokenizer( + example[dataset_text_field] if not do_formatting_func else formatting_func(example), + truncation = do_truncation, + max_length = max_seq_length, + return_token_type_ids = _needs_token_type_ids, + add_special_tokens = add_special_tokens, + ) + pass + + if not isinstance(dataset, IterableDataset): + import multiprocessing as _mp + dataset_num_proc = getattr(args, "dataset_num_proc", None) + if dataset_num_proc is None: + if _mp.get_start_method() != 'fork': + dataset_num_proc = None + else: + import psutil + dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64) + memory_gb_left = psutil.virtual_memory().available / (1024**3) + if memory_gb_left <= 2: + dataset_num_proc = 1 + else: + dataset_num_proc = min(dataset_num_proc, int(memory_gb_left)) + map_kwargs["num_proc"] = dataset_num_proc + else: + map_kwargs["batch_size"] = dataset._ex_iterable.batch_size + + if do_prompt_completion: + _eos_token = getattr(tokenizer, 'eos_token', None) + + def _tokenize_pc(example): + if _is_conv: + prompt_ids = processing_class.apply_chat_template( + example["prompt"], tokenize=True, + add_generation_prompt=True, return_dict=False, + tools=example.get("tools"), + **(example.get("chat_template_kwargs") or {}), + ) + if prompt_ids and isinstance(prompt_ids[0], list): + prompt_ids = prompt_ids[0] + pc_processed = processing_class.apply_chat_template( + example["prompt"] + example["completion"], + return_dict=True, tokenize=True, + tools=example.get("tools"), + **(example.get("chat_template_kwargs") or {}), + ) + if isinstance(pc_processed.get("input_ids", [None])[0], list): + pc_processed = {k: v[0] for k, v in pc_processed.items()} + pc_ids = pc_processed["input_ids"] + else: + _completion = example["completion"] + if _eos_token and not _completion.endswith(_eos_token): + _completion = _completion + _eos_token + prompt_ids = tokenizer( + example["prompt"], add_special_tokens=add_special_tokens, + )["input_ids"] + pc_ids = tokenizer( + example["prompt"] + _completion, + add_special_tokens=add_special_tokens, + )["input_ids"] + if do_truncation and max_seq_length > 0: + pc_ids = pc_ids[:max_seq_length] + n_prompt = min(len(prompt_ids), len(pc_ids)) + completion_mask = [0] * n_prompt + [1] * (len(pc_ids) - n_prompt) + result = {"input_ids": pc_ids, "completion_mask": completion_mask} + if _needs_token_type_ids: + result["token_type_ids"] = [0] * len(pc_ids) + return result + + if use_desc: + map_kwargs["desc"] = 'Unsloth: Tokenizing ["prompt"+"completion"]' + import warnings as _w + with _w.catch_warnings(): + _w.filterwarnings("ignore", message=".*couldn't be hashed properly.*") + dataset = dataset.map( + _tokenize_pc, batched=False, + remove_columns=list(column_names), **map_kwargs, + ) + else: + if use_desc: map_kwargs["desc"] = f'Unsloth: Tokenizing ["{dataset_text_field}"]' + import warnings as _w + with _w.catch_warnings(): + _w.filterwarnings("ignore", message=".*couldn't be hashed properly.*") + dataset = dataset.map(_tokenize, batched = True, remove_columns = list(column_names), **map_kwargs) + + # VLMs need .pad; switch the data collator + if is_vlm and not hasattr(processing_class, "pad") and not do_prompt_completion: + data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) + self.data_collator = data_collator + pass + pass + if packing: + # Use TRL's pack_dataset if available + try: + pack_dataset + except: + print("Unsloth: Hugging Face's packing is currently buggy - we're disabling it for now!") + return dataset + + if max_seq_length == 0: + raise ValueError("When packing is enabled, `max_seq_length` can't be `None`.") + + if use_desc: map_kwargs["desc"] = f"Unsloth: Packing {dataset_name} dataset" + dataset = pack_dataset( + dataset.select_columns(used_column_names), + max_seq_length, + getattr(args, "packing_strategy", "bfd"), + map_kwargs, + ) + pass + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids", + # "attention_mask" and "labels"). Dataset preparation also produces a "seq_lengths" column (for packing / + # padding-free), so we override the default signature columns to keep it alongside the model inputs. + if self._signature_columns is None: + if self._is_vision_dataset: + self._signature_columns = ["messages", "prompt", "completion", "image", "images"] + else: + self._signature_columns = ["input_ids", "labels", "seq_lengths"] + + def _reject_skip_prepare_without_labels(self, datasets: dict[str, Dataset], data_collator) -> None: + # This guard may look defensive, but it covers a behavior change introduced when label building moved from + # the collator to dataset preparation: the collator used to consume the mask columns directly, so a + # skipped-preparation dataset carrying masks trained correctly. Now labels are built during preparation, which + # is skipped here, and the collator ignores the mask columns. Without a "labels" column, such a dataset would + # silently optimize the loss over the full sequence, so we fail loudly instead. Checked both at init and in + # `evaluate`, since a dataset passed directly to `evaluate` also skips preparation. + if not ( + self._skip_prepare_dataset + and not self._is_vision_dataset + and isinstance(data_collator, DataCollatorForLanguageModeling) + ): + return + for name, dataset in datasets.items(): + cols = get_dataset_column_names(dataset) + if "labels" not in cols and ("completion_mask" in cols or "assistant_masks" in cols): + raise ValueError( + f"The {name} dataset has mask columns but no 'labels', and `skip_prepare_dataset=True` skips " + "label building, so it would train on the full sequence. Add a 'labels' column (-100 for " + "non-loss tokens) or drop `skip_prepare_dataset`." + ) + + def evaluate( + self, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> dict[str, float]: + # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way + # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer (language modeling, + # prompt-completion, etc.). `_prepare_dataset` is idempotent: it skips datasets that are already tokenized. A + # `str` selects a dataset that was already prepared at init time, so it's left untouched. + if not self._skip_prepare_dataset and eval_dataset is not None and not isinstance(eval_dataset, str): + packing = self.args.packing if self.args.eval_packing is None else self.args.eval_packing + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset( + dataset, self.processing_class, self.args, packing, self._formatting_func, key + ) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset( + eval_dataset, self.processing_class, self.args, packing, self._formatting_func, "eval" + ) + eval_datasets = ( + eval_dataset + if isinstance(eval_dataset, dict) + else {"eval": eval_dataset} + if eval_dataset is not None and not isinstance(eval_dataset, str) + else {} + ) + self._reject_skip_prepare_without_labels(eval_datasets, self.data_collator) + return super().evaluate( + eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix + ) + + def compute_loss( + self, + model, + inputs, + return_outputs = False, + num_items_in_batch = None, + ): + outputs = super().compute_loss( + model, + inputs, + return_outputs = return_outputs, + num_items_in_batch = num_items_in_batch, + ) + return outputs + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): + # Preserve the eval loop intent so compute_loss can decide whether logits are needed. + inputs["_prediction_loss_only"] = prediction_loss_only + return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys) + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) +class UnslothSFTTrainer(_UnslothSFTTrainer): + """ + + Trainer for Supervised Fine-Tuning (SFT) method. + + This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + >>> from trl import SFTTrainer + >>> from datasets import load_dataset + + >>> dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]") + + >>> trainer = SFTTrainer( + ... model="Qwen/Qwen2.5-0.5B-Instruct", + ... train_dataset=dataset, + ... ) + >>> trainer.train() + ``` + + Args: + model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `.from_pretrained` (where `` is derived from the model + config) with the keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + - A [`~peft.PeftModel`] object. Only causal language models are supported. + args ([`SFTConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator ([`~transformers.DataCollator`], *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~trainer.sft_trainer.DataCollatorForLanguageModeling`] if the model is a language model + and [`~trainer.sft_trainer.DataCollatorForVisionLanguageModeling`] if the model is a vision-language model. + Custom collators must truncate sequences before padding; the trainer does not apply post-collation + truncation. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. This trainer supports both [language modeling](#language-modeling) type and + [prompt-completion](#prompt-completion) type. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + + The trainer also supports processed datasets (tokenized) as long as they contain an `input_ids` field. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoProcessor.from_pretrained`]. A padding token, `tokenizer.pad_token`, must be set. + If the processing class has not set a padding token, `tokenizer.eos_token` will be used as the default. + compute_loss_func (`Callable`, *optional*): + A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated + batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default [loss + function](https://github.com/huggingface/transformers/blob/052e652d6d53c2b26ffde87e039b723949a53493/src/transformers/trainer.py#L3618) + used by [`Trainer`]. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a + [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing + [`SFTConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a boolean + `compute_result` argument. This will be triggered after the last eval batch to signal that the function + needs to calculate and return the global summary statistics rather than accumulating the batch-level + statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before + initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + formatting_func (`Callable`, *optional*): + Formatting function applied to the dataset before tokenization. Applying the formatting function explicitly + converts the dataset into a [language modeling](#language-modeling) type. + + """ + def __init__( + self, + model, + args = None, + data_collator = None, + train_dataset = None, + eval_dataset = None, + processing_class = None, + compute_loss_func = None, + compute_metrics = None, + callbacks = None, + optimizer_cls_and_kwargs = None, + preprocess_logits_for_metrics = None, + peft_config = None, + formatting_func = None, + **kwargs + ): + if args is None: args = UnslothSFTConfig() + use_bf16 = getattr(args, 'bf16', False) + if type(use_bf16) is not bool: use_bf16 = False + use_fp16 = getattr(args, 'fp16', False) + if type(use_fp16) is not bool: use_fp16 = False + force_float32 = False + full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' + if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'): + print('Unsloth: Switching to float32 training since model cannot work with float16') + force_float32 = True + mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') + dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) + if dtype is None: dtype = model.get_input_embeddings().weight.dtype + from unsloth_zoo.utils import _get_dtype + dtype = _get_dtype(dtype) + float16 = dtype == torch.float16 + if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') + if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') + if force_float32: + # Forced float32 training + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': + # Mixed precision training + args.fp16 = float16 + args.bf16 = not float16 + os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' + # args.mixed_precision is a new argument which needs to be set now + elif mixed_precision_dtype == 'bfloat16': + # Both False since bfloat16 full finetuning doesn't do any autocasting. + args.fp16 = False + args.bf16 = False + os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' + if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' + # args.mixed_precision is a new argument which needs to be set now + + if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': + args.eval_strategy = 'steps' + if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 + ga_steps = getattr(args, 'gradient_accumulation_steps', None) + if ga_steps is not None and ga_steps > 1: + from transformers import __version__ as transformers_version + if Version(transformers_version) <= Version('4.45.2'): + print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' + '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') + if getattr(args, 'eval_strategy', 'no') != 'no': + eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) + if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size + if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps + fp16_full_eval = getattr(args, 'fp16_full_eval', False) + if type(fp16_full_eval) is not bool: fp16_full_eval = False + bf16_full_eval = getattr(args, 'bf16_full_eval', False) + if type(bf16_full_eval) is not bool: bf16_full_eval = False + if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True + if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False + if force_float32: + args.bf16_full_eval = False + args.fp16_full_eval = False + elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': + args.bf16_full_eval = True + args.fp16_full_eval = False + elif not bf16_full_eval and not fp16_full_eval: + args.bf16_full_eval = args.bf16 + args.fp16_full_eval = args.fp16 + _output_logits = False + if locals().get('compute_metrics', None) is not None: _output_logits = True + if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True + if _output_logits: + os.environ['UNSLOTH_RETURN_LOGITS'] = '1' + if model is not None: + _warnings_issued = getattr(model, 'warnings_issued', None) + if _warnings_issued is None: + model.warnings_issued = {} + elif not isinstance(_warnings_issued, dict): + try: + model.warnings_issued = dict(_warnings_issued) + except Exception: + model.warnings_issued = {} + if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): + pass + else: + model_max_seq_length = getattr(model, 'max_seq_length', None) + args_max_seq_length = getattr(args, 'max_seq_length', None) + if args_max_seq_length is None and model_max_seq_length is not None: + max_seq_length = model.max_seq_length + if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length + elif args_max_seq_length is not None and model_max_seq_length is not None: + if args_max_seq_length > model_max_seq_length: + print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' + 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') + args.max_seq_length = model_max_seq_length + if 'max_length' not in locals() and not hasattr(args, 'max_length'): + pass + else: + if hasattr(args, 'max_seq_length') and args.max_seq_length is not None and args.max_seq_length > 0: + if hasattr(args, 'max_length'): + args.max_length = args.max_seq_length + max_length = args.max_length + else: + model_max_length = getattr(model, 'max_seq_length', None) + if model_max_length is None: model_max_length = getattr(model, 'max_length', None) + if model_max_length is not None: + args.max_length = model_max_length + max_length = args.max_length + elif hasattr(args, 'max_length') and args.max_length is not None: + max_length = args.max_length + # if we are here, then we are in a weird case where max_length is set but max_seq_length is not set + setattr(model, 'max_seq_length', max_length) + else: + print('Unsloth: We did not find `max_seq_length` or `max_length` in the model or args. We will set it to 1024.') + args.max_length = 1024 + if model is not None and hasattr(model, 'for_training'): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' + if 'processing_class' in locals(): + if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' + if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' + __tokenizer = processing_class if 'processing_class' in locals() else tokenizer + from unsloth_zoo.vision_utils import UnslothVisionDataCollator + if not isinstance(data_collator, UnslothVisionDataCollator): + if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names: + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names: + data_collator = DataCollatorForSeq2Seq( + __tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + else: + if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False + if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' + if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} + if not isinstance(data_collator, UnslothVisionDataCollator): + if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): + if isinstance(data_collator, DataCollatorForSeq2Seq): + data_collator = DataCollatorForSeq2Seq( + __tokenizer.tokenizer, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): + data_collator = TransformersDataCollatorForLanguageModeling( + __tokenizer.tokenizer, + mlm = False, + mlm_probability = 0.0, + pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), + ) + other_metrics = [] + + from unsloth_zoo.logging_utils import PatchRLStatistics + PatchRLStatistics('sft_trainer', other_metrics) + IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', '').split('\n') + from unsloth_zoo.tokenizer_utils import fix_untrained_tokens + from unsloth_zoo.training_utils import fix_zero_training_loss + if 'tokenizer' not in locals(): tokenizer = processing_class + fix_untrained_tokens(model, tokenizer, train_dataset, IGNORED_TOKENIZER_NAMES, eps = 1e-16) + fix_zero_training_loss(model, tokenizer, train_dataset) + + # [TODO] Fix up DataParallel multiplying batch sizes + # [TODO] DDP works, but DP seems to not work? [TODO] + if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: + if getattr(args, "_n_gpu", 1) != 1: + args._n_gpu = 1 + if "model" in locals() and hasattr(model, "for_training"): + model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + super().__init__( + model = model, + args = args, + data_collator = data_collator, + train_dataset = train_dataset, + eval_dataset = eval_dataset, + processing_class = processing_class, + compute_loss_func = compute_loss_func, + compute_metrics = compute_metrics, + callbacks = callbacks, + optimizer_cls_and_kwargs = optimizer_cls_and_kwargs, + preprocess_logits_for_metrics = preprocess_logits_for_metrics, + peft_config = peft_config, + formatting_func = formatting_func,**kwargs) + if "model" in locals() and hasattr(model, "for_inference"): + model.for_inference() + if hasattr(self, 'neftune_hook_handle'): + self.neftune_hook_handle.remove() + if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle + if getattr(args, 'neftune_noise_alpha', None) is not None: + model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha + pass + if hasattr(self, 'accelerator'): + scaler = self.accelerator.scaler + current_model = model + while hasattr(current_model, 'model'): + current_model.accelerator_scaler = scaler + current_model = current_model.model + current_model.accelerator_scaler = scaler + pass + if hasattr(self, 'train'): + self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) + pass + if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): + _vllm_tok = self.llm.get_tokenizer() + _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) + if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: + _vllm_tok.chat_template = _pc.chat_template + pass + +pass + + +if hasattr(logger, "addFilter"): + import logging + class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not (self.text in x.getMessage()) + pass + logger.addFilter(HideLoggingMessage("`use_cache=True`")) + diff --git a/unsloth_compiled_cache/__pycache__/AqlmLoraLinear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/AqlmLoraLinear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d279d334a7878ab03064fdd26d147d54bdbe0be Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/AqlmLoraLinear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/AwqLoraLinear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/AwqLoraLinear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0882da505040baf1c6003747328f13447ef4192 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/AwqLoraLinear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/BatchNorm1d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/BatchNorm1d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a48dcb7dcf0410bb7e68801e93f0c1e2041355d Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/BatchNorm1d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/BatchNorm2d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/BatchNorm2d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83035fae8b0ede69d4e5588b7527cda04680a921 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/BatchNorm2d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/BatchNorm3d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/BatchNorm3d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f54fad628d84874a9064dcb2b7d6c049c614379 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/BatchNorm3d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/BlockDiagonalLinear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/BlockDiagonalLinear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9470565be8ea191177b2347d12e6d22b3ff213d7 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/BlockDiagonalLinear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/Conv1d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/Conv1d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c4e2eb9a63093fe25418f653ee65afbe659b9c5 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/Conv1d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/Conv2d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/Conv2d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..798cab6291e4434aa13c5613650a2123df969544 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/Conv2d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/Conv3d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/Conv3d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7578a3907eaf51aed12b637c10c26a1ea1a24606 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/Conv3d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/ConvTranspose1d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/ConvTranspose1d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0901ac3514acca5931c20397b74a7be5d6230c88 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/ConvTranspose1d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/ConvTranspose2d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/ConvTranspose2d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03c5ecd5275862a8f81c24603b401d0e4ebaf4cd Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/ConvTranspose2d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/ConvTranspose3d.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/ConvTranspose3d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f3a420e39765b2809dd205f319728eb51c03085b Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/ConvTranspose3d.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/GPTQLoraLinear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/GPTQLoraLinear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63baf5d5bba1b2026e23468d4f21975fe19c7207 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/GPTQLoraLinear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/GroupNorm.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/GroupNorm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..385f27e6f6f6971b00a00727f219c8da84150f3b Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/GroupNorm.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/LayerNorm.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/LayerNorm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..150b9c9fb4fdb5bff49823b1a29d38f77b4c38f3 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/LayerNorm.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/Linear4bit_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/Linear4bit_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25f2a4ebb62a98179c5ad22b78386df07071cc8d Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/Linear4bit_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/Linear8bitLt_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/Linear8bitLt_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d1e700beab9d1f9e2c17d8a2b0e26c8b292a46bb Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/Linear8bitLt_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/Linear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/Linear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a6bd5d39696f4d7e89f6e5aafe501de1f6ea4ae Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/Linear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/LoraParallelLinear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/LoraParallelLinear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22320a5f7c8ff4b9d008581f3c56920534b0667d Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/LoraParallelLinear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/RMSNorm.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/RMSNorm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..771ed90b5665438d1a17690b3c1418a6729c504f Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/RMSNorm.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/TeLinear_peft_forward.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/TeLinear_peft_forward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e14fffa2dce9cd8c123df2e838dc34a31b800ce5 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/TeLinear_peft_forward.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/UnslothDPOTrainer.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/UnslothDPOTrainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d0fa715aa4a226773e6d9658ff5f9760e1c10f0 --- /dev/null +++ b/unsloth_compiled_cache/__pycache__/UnslothDPOTrainer.cpython-312.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6ed89dedd21fef6dd7c419f02bda55e7ba7d8f8905133d9fb7b8da9937b00ad0 +size 113852 diff --git a/unsloth_compiled_cache/__pycache__/UnslothGRPOTrainer.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/UnslothGRPOTrainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..527d124698d1cb1df86d440f1e63ba00c96b4826 --- /dev/null +++ b/unsloth_compiled_cache/__pycache__/UnslothGRPOTrainer.cpython-312.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2af75462e74d3fed374f1147d89e19dcba68bfbab16efcb1101c9b20c339b4ff +size 261874 diff --git a/unsloth_compiled_cache/__pycache__/UnslothKTOTrainer.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/UnslothKTOTrainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4d8e06f79ad026c497709d65c96e8308575d7ac --- /dev/null +++ b/unsloth_compiled_cache/__pycache__/UnslothKTOTrainer.cpython-312.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7767f679b6f8f6037aba74efae0dad1b57ef26b951c7394f58ebf36efabad999 +size 100160 diff --git a/unsloth_compiled_cache/__pycache__/UnslothRLOOTrainer.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/UnslothRLOOTrainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fcabf88ca1bf076b24b432b4c478edd0c1c66da3 --- /dev/null +++ b/unsloth_compiled_cache/__pycache__/UnslothRLOOTrainer.cpython-312.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4a304fa90948ec1e4aa2d3d68ae3546ed9e90174fd9cb810cf02f96754e80fa +size 132490 diff --git a/unsloth_compiled_cache/__pycache__/UnslothRewardTrainer.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/UnslothRewardTrainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5722b79dea506c78737893b12488d456de0e3323 Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/UnslothRewardTrainer.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/UnslothSFTTrainer.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/UnslothSFTTrainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9adc66be5762a13a976535b92a71b35b2900665c Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/UnslothSFTTrainer.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/moe_utils.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/moe_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1819e8d08eb78f749865fa97a51afab78b7e6e7c Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/moe_utils.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/unsloth_compiled_module_gemma4.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/unsloth_compiled_module_gemma4.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c6eed9ce645e4b994a4607055b530c109000d0f Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/unsloth_compiled_module_gemma4.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/__pycache__/unsloth_compiled_module_siglip.cpython-312.pyc b/unsloth_compiled_cache/__pycache__/unsloth_compiled_module_siglip.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa4ef2976fb4ab708f28e4a28483a095f6ade72e Binary files /dev/null and b/unsloth_compiled_cache/__pycache__/unsloth_compiled_module_siglip.cpython-312.pyc differ diff --git a/unsloth_compiled_cache/moe_utils.py b/unsloth_compiled_cache/moe_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7f29dcfeb2a434489a38d5bbf84458af9ab9bd4f --- /dev/null +++ b/unsloth_compiled_cache/moe_utils.py @@ -0,0 +1,1587 @@ +# Unsloth Zoo - Utilities for Unsloth +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +import torch +import torch.nn.functional as F +import os +import shutil +import sys +import importlib +import importlib.util +from typing import Optional, Tuple +from torch.autograd import Function +from unsloth_zoo.mlx import is_mlx_available + +UNSLOTH_COMPILE_LOCATION = os.environ.get( + "UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache" +) + +try: + import bitsandbytes as bnb + from bitsandbytes.nn import Params4bit + HAS_BNB = True +except ImportError: + HAS_BNB = False + Params4bit = None + + +def _get_compile_location() -> str: + return os.path.abspath( + os.environ.get("UNSLOTH_COMPILE_LOCATION", UNSLOTH_COMPILE_LOCATION) + ) + + +def _log_info(message: str): + if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1": + print(message) + + +def install_to_cache(source_path, destination_filename=None): + """Copy a file into unsloth_compiled_cache so compiled modules can use it.""" + compile_location = _get_compile_location() + if not os.path.exists(compile_location): + try: + os.makedirs(compile_location) + except: + pass + + current_file = os.path.abspath(source_path) + if destination_filename is None: + destination_filename = os.path.basename(current_file) + + destination = os.path.abspath(os.path.join(compile_location, destination_filename)) + + if current_file != destination: + try: + shutil.copy(current_file, destination) + except Exception: + pass + + +install_to_cache(__file__, "moe_utils.py") + +_CACHED_FORWARD_MOE_BACKEND = None +_CACHED_MOE_UTILS_MODULE = None + + +def _load_cached_moe_utils_module(): + global _CACHED_MOE_UTILS_MODULE + + cache_file = os.path.abspath(os.path.join(_get_compile_location(), "moe_utils.py")) + current_file = os.path.abspath(__file__) + if not os.path.isfile(cache_file) or cache_file == current_file: + return None + + try: + module_name = "unsloth_cached_moe_utils" + module = sys.modules.get(module_name, None) + if module is not None and os.path.abspath(getattr(module, "__file__", "")) == cache_file: + _CACHED_MOE_UTILS_MODULE = module + return module + + spec = importlib.util.spec_from_file_location(module_name, cache_file) + if spec is None or spec.loader is None: + return None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + _CACHED_MOE_UTILS_MODULE = module + return module + except Exception: + return None + + +def get_forward_moe_backend(): + """Resolve forward_moe_backend from the compiled cache copy, else the local def.""" + global _CACHED_FORWARD_MOE_BACKEND + module = _load_cached_moe_utils_module() + if module is not None and hasattr(module, "forward_moe_backend"): + _CACHED_FORWARD_MOE_BACKEND = module.forward_moe_backend + return _CACHED_FORWARD_MOE_BACKEND + + _CACHED_FORWARD_MOE_BACKEND = forward_moe_backend + return _CACHED_FORWARD_MOE_BACKEND + +# Grouped MM wrapper around torch._grouped_mm; native backward works correctly. + + +def _grouped_mm_with_backward_fix( + inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor +) -> torch.Tensor: + """Grouped matmul via torch._grouped_mm with contiguous inputs. + + Some low-rank LoRA weights are contiguous but still have row strides below + the kernel alignment requirement, so keep a narrow fallback for those cases. + """ + inputs = inputs.contiguous() + weight = weight.contiguous() + try: + return torch._grouped_mm(inputs, weight, offs=offsets) + except RuntimeError as exc: + message = str(exc) + if "strides should be multiple of 16 bytes" not in message: + raise + return _manual_grouped_mm(inputs, weight, offsets) + + +def _manual_grouped_mm( + inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor +) -> torch.Tensor: + """Differentiable grouped matmul fallback for torch._grouped_mm alignment gaps.""" + outputs = [] + start = 0 + for expert_idx, end in enumerate(offsets.detach().cpu().tolist()): + if start < end: + outputs.append(torch.matmul(inputs[start:end], weight[expert_idx])) + start = end + if outputs: + return torch.cat(outputs, dim=0) + return inputs.new_empty((0, weight.shape[-1])) + + +_GROUPED_GEMM_AVAILABLE = None +_TORCH_GROUPED_MM_AVAILABLE = hasattr(torch, "_grouped_mm") + +# GPU support for torch._grouped_mm, verified via runtime probe. +_TORCH_GROUPED_MM_SUPPORTED = None + + +def _check_torch_grouped_mm_supported(): + """Check torch._grouped_mm support on the current GPU; a runtime probe is the only reliable check.""" + global _TORCH_GROUPED_MM_SUPPORTED + if _TORCH_GROUPED_MM_SUPPORTED is not None: return _TORCH_GROUPED_MM_SUPPORTED + + if not _TORCH_GROUPED_MM_AVAILABLE: + _TORCH_GROUPED_MM_SUPPORTED = False + return False + + if not torch.cuda.is_available(): + _TORCH_GROUPED_MM_SUPPORTED = False + return False + + try: + # Dummy call verifies real support (symbol may exist but hardware unsupported, e.g. < H100). + device = torch.cuda.current_device() + dtype = torch.float16 + + # 1 expert, 1 token, dim 8 (safe alignment). + x = torch.ones((1, 8), device=device, dtype=dtype) + w = torch.ones((1, 8, 8), device=device, dtype=dtype) + offs = torch.tensor([1], device=device, dtype=torch.int32) + + torch._grouped_mm(x, w, offs=offs) + del x, w, offs + _TORCH_GROUPED_MM_SUPPORTED = True + except Exception: + _TORCH_GROUPED_MM_SUPPORTED = False + + return _TORCH_GROUPED_MM_SUPPORTED + + +_TRITON_ALLOCATOR_INITIALIZED = False +_PERSISTENT_BUFFER = None +_original_peft_get_peft_model = None + + +def _init_triton_allocator(): + """Initialize a persistent Triton allocator to avoid per-call allocation overhead.""" + global _TRITON_ALLOCATOR_INITIALIZED, _PERSISTENT_BUFFER + if _TRITON_ALLOCATOR_INITIALIZED: return + + try: + import triton + + # Persistent buffer that grows as needed, avoiding per-kernel allocations. + def persistent_alloc_fn(size: int, alignment: int, stream): + global _PERSISTENT_BUFFER + # Round up to nearest 128 bytes for alignment / fewer reallocations. + rounded_size = ((size + 128 - 1) // 128) * 128 + + if ( + _PERSISTENT_BUFFER is None + or _PERSISTENT_BUFFER.numel() * _PERSISTENT_BUFFER.element_size() + < rounded_size + ): + # 10% headroom; uint8 for raw byte storage. + _PERSISTENT_BUFFER = torch.empty( + int(rounded_size * 1.1), device="cuda", dtype=torch.uint8 + ) + _PERSISTENT_BUFFER.__hibernate__ = {"type": "ignore"} + return _PERSISTENT_BUFFER + + triton.set_allocator(persistent_alloc_fn) + triton._unsloth_allocator_set = True + _TRITON_ALLOCATOR_INITIALIZED = True + except Exception: + pass + + +def _check_grouped_gemm_available(): + """Check if Unsloth grouped GEMM kernels are available.""" + if os.environ.get("UNSLOTH_DISABLE_MOE_TRITON", "0") == "1": return False + if is_mlx_available(): return False + + global _GROUPED_GEMM_AVAILABLE + if _GROUPED_GEMM_AVAILABLE is not None: return _GROUPED_GEMM_AVAILABLE + + try: + from unsloth.kernels.moe.grouped_gemm.interface import grouped_gemm, supports_tma + _GROUPED_GEMM_AVAILABLE = True + _init_triton_allocator() + except (ImportError, ModuleNotFoundError): + _GROUPED_GEMM_AVAILABLE = False + return _GROUPED_GEMM_AVAILABLE + + +from functools import lru_cache, wraps + + +@lru_cache(maxsize=1) +def select_moe_backend(): + """Select MoE backend from UNSLOTH_MOE_BACKEND + availability. + + Choices: "grouped_mm", "unsloth_triton", "native_torch" (default "grouped_mm"). + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + requested = os.environ.get("UNSLOTH_MOE_BACKEND") + if requested: + if requested == "grouped_mm" and _check_torch_grouped_mm_supported(): + return "grouped_mm" + if requested == "unsloth_triton" and _check_grouped_gemm_available(): + return "unsloth_triton" + if requested == "native_torch": + return "native_torch" + _log_info(f"Unsloth: '{requested}' backend requested but is not available. Falling back to next available.") + + if _check_torch_grouped_mm_supported(): + _log_info("Unsloth: Using MoE backend 'grouped_mm'") + return "grouped_mm" + if _check_grouped_gemm_available(): + _log_info("Unsloth: Using MoE backend 'unsloth_triton'") + return "unsloth_triton" + return "native_torch" + + +def swap_moe_weights_for_call(experts_module, gate_up_proj, down_proj, forward_fn, *args): + """Temporarily install dequantized weights for one forward call, then restore. + + Uses object.__setattr__ to bypass nn.Module Parameter (de)registration + (re-registers hooks, unnecessary for read-only temp tensors). Used by the + FP8 and bnb4bit MoE dispatchers. + """ + original_gate_up = experts_module.gate_up_proj + original_down = experts_module.down_proj + object.__setattr__(experts_module, "gate_up_proj", gate_up_proj) + object.__setattr__(experts_module, "down_proj", down_proj) + try: + return forward_fn(experts_module, *args) + finally: + object.__setattr__(experts_module, "gate_up_proj", original_gate_up) + object.__setattr__(experts_module, "down_proj", original_down) + + +def forward_moe_backend( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Dispatch MoE forward to the selected backend (keeps model-specific patches minimal).""" + # This Unsloth Zoo code section is licensed under AGPL3 + + # Absolute imports: this function is also copied into + # unsloth_compiled_cache/moe_utils.py where relative imports of sibling + # helpers don't resolve (only the dispatcher is copied). + # Keep `except ImportError` around ONLY the import; runtime errors in the + # bnb4bit/fp8 path must propagate, not fall through to a crashing backend. + _moe_uses_bnb4bit_expert_weights = forward_moe_backend_bnb4bit = None + try: + from unsloth_zoo.temporary_patches.moe_utils_bnb4bit import ( + _moe_uses_bnb4bit_expert_weights, + forward_moe_backend_bnb4bit, + ) + except ImportError: + pass + if _moe_uses_bnb4bit_expert_weights is not None and _moe_uses_bnb4bit_expert_weights(self): + result = forward_moe_backend_bnb4bit(self, hidden_states, top_k_index, top_k_weights) + if result is not None: + return result + + _moe_uses_fp8_expert_weights = forward_moe_backend_fp8 = None + try: + from unsloth_zoo.temporary_patches.moe_utils_fp8 import ( + _moe_uses_fp8_expert_weights, + forward_moe_backend_fp8, + ) + except ImportError: + pass + if _moe_uses_fp8_expert_weights is not None and _moe_uses_fp8_expert_weights(self): + return forward_moe_backend_fp8(self, hidden_states, top_k_index, top_k_weights) + + backend = select_moe_backend() + if backend == "grouped_mm": + return forward_native_grouped_mm(self, hidden_states, top_k_index, top_k_weights) + if backend == "unsloth_triton": + return forward_triton_grouped_gemm(self, hidden_states, top_k_index, top_k_weights) + return forward_native_moe_loop(self, hidden_states, top_k_index, top_k_weights) + + +@torch.no_grad() +def _get_routing_indices(selected_experts, num_experts): + """Compute token->expert mapping for grouped GEMM. + + Returns (token_counts_by_expert (num_experts,), gather_indices (total_tokens,)). + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + flat_experts = selected_experts.view(-1) + + # bincount avoids histc's float conversion overhead. + token_counts_by_expert = torch.bincount(flat_experts, minlength=num_experts).to(torch.int32) + + # stable=True preserves order within each expert. + gather_indices = flat_experts.argsort(stable=True) + + return token_counts_by_expert, gather_indices + + +def _silu_and_mul(x): + """Fused SiLU + element-wise multiply for gate/up projections.""" + gate, up = x.chunk(2, dim=-1) + return F.silu(gate) * up + + +# Separated LoRA helpers. + + +def _has_lora_adapters(param) -> bool: + """Check for active LoRA adapters (PEFT ParamWrapper).""" + if not hasattr(param, "lora_A") or not hasattr(param, "lora_B"): + return False + if hasattr(param, "disable_adapters") and param.disable_adapters: + return False + if hasattr(param, "merged") and param.merged: + return False + return len(param.lora_A) > 0 + + +def _canonical_lora_weights_for_grouped_mm( + weight_A: torch.Tensor, + weight_B: torch.Tensor, + num_experts: int, + rank_per_expert: int, + dim_A: int, + dim_B: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + first_weight = weight_A.view(num_experts, rank_per_expert, dim_A) + first_weight = first_weight.permute(0, 2, 1).contiguous() + second_weight = weight_B.view(dim_B, num_experts, rank_per_expert) + second_weight = second_weight.permute(1, 2, 0).contiguous() + return first_weight, second_weight + + +def _reversed_lora_weights_for_grouped_mm( + weight_A: torch.Tensor, + weight_B: torch.Tensor, + num_experts: int, + rank_per_expert: int, + dim_A: int, + dim_B: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + first_weight = weight_B.view(dim_B, num_experts, rank_per_expert) + first_weight = first_weight.permute(1, 0, 2).contiguous() + second_weight = weight_A.view(num_experts, rank_per_expert, dim_A).contiguous() + return first_weight, second_weight + + +def _get_param_shape_from_module(module, parameter_name): + if module is None or parameter_name is None or not hasattr(module, parameter_name): + return None + param = getattr(module, parameter_name) + if hasattr(param, "get_param"): + param = param.get_param() + elif hasattr(param, "weight"): + param = param.weight + return tuple(param.shape) + + +def _get_moe_lora_io_dims(wrapper, experts_module=None): + base = None + if wrapper is not None and hasattr(wrapper, "get_base_layer"): + base = wrapper.get_base_layer() + if experts_module is None: + experts_module = base + if experts_module is None: + experts_module = getattr(wrapper, "base_layer", None) + + parameter_name = getattr(wrapper, "parameter_name", None) + source = experts_module if experts_module is not None else base + if source is None: + return None, None + _set_gpt_oss_grouped_mm_format_on_experts(source) + + shape = _get_param_shape_from_module(source, parameter_name) + if shape is not None and len(shape) >= 3: + grouped_mm_format = bool(getattr(source, "_unsloth_grouped_mm_format", False)) + if grouped_mm_format: + return shape[-2], shape[-1] + return shape[-1], shape[-2] + + hidden_dim = getattr(source, "hidden_dim", None) + intermediate_dim = getattr(source, "intermediate_dim", None) + if hidden_dim is None or intermediate_dim is None: + return None, None + if parameter_name == "gate_up_proj": + return hidden_dim, 2 * intermediate_dim + if parameter_name == "down_proj": + return intermediate_dim, hidden_dim + return None, None + + +def extract_moe_lora_weights_for_grouped_mm( + wrapper, + weight_A: torch.Tensor, + weight_B: torch.Tensor, + scaling, + num_experts: int, + *, + experts_module=None, + input_dim=None, + output_dim=None, + model_name: str = "MoE", + enable_logging: bool = None, + logger_obj=None, +) -> Tuple[torch.Tensor, torch.Tensor, float, int]: + total_rank = weight_A.shape[0] + rank_per_expert = total_rank // num_experts + dim_A = weight_A.shape[1] + dim_B = weight_B.shape[0] + + if num_experts <= 1: + return weight_A.T, weight_B.T, scaling, num_experts + + if input_dim is None or output_dim is None: + inferred_input_dim, inferred_output_dim = _get_moe_lora_io_dims( + wrapper, experts_module=experts_module, + ) + if input_dim is None: + input_dim = inferred_input_dim + if output_dim is None: + output_dim = inferred_output_dim + + canonical_match = ( + input_dim is not None + and output_dim is not None + and dim_A == input_dim + and dim_B == output_dim + ) + reversed_match = ( + input_dim is not None + and output_dim is not None + and dim_A == output_dim + and dim_B == input_dim + ) + + if canonical_match and reversed_match: + if bool(getattr(wrapper, "_did_swap_in_out_features", False)): + first_weight, second_weight = _reversed_lora_weights_for_grouped_mm( + weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, + ) + else: + first_weight, second_weight = _canonical_lora_weights_for_grouped_mm( + weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, + ) + return first_weight, second_weight, scaling, num_experts + + if canonical_match: + first_weight, second_weight = _canonical_lora_weights_for_grouped_mm( + weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, + ) + return first_weight, second_weight, scaling, num_experts + + if reversed_match: + first_weight, second_weight = _reversed_lora_weights_for_grouped_mm( + weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, + ) + return first_weight, second_weight, scaling, num_experts + + if logger_obj is not None: + if enable_logging is None: + enable_logging = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" + if enable_logging and (input_dim is not None or output_dim is not None): + logger_obj.warning( + f"Unsloth: {model_name} LoRA extractor could not match either layout " + f"(weight_A={tuple(weight_A.shape)}, weight_B={tuple(weight_B.shape)}, " + f"expected input_dim={input_dim}, output_dim={output_dim}, " + f"num_experts={num_experts}). Falling back to canonical layout. " + "If this is a new PEFT version, the LoRA delta may be wrong." + ) + + first_weight, second_weight = _canonical_lora_weights_for_grouped_mm( + weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, + ) + return first_weight, second_weight, scaling, num_experts + + +def _extract_lora_from_wrapper( + wrapper, adapter_name: str = "default", experts_module=None +) -> Optional[Tuple[torch.Tensor, torch.Tensor, float, int]]: + """Extract LoRA weights from a PEFT ParamWrapper for MoE separated grouped_mm. + + PEFT 3D ParamWrapper gives lora_A: (E*R, in_dim), lora_B: (out_dim, E*R); + reshaped to first_weight (E, in_dim, R), second_weight (E, R, out_dim) so + delta = X @ first @ second. Handles both standard (E, out, in) Qwen3-MoE and + transposed (E, in, out) Qwen3-VL-MoE base weight layouts. + + Returns (first_weight, second_weight, scaling, num_experts) or None. + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + try: + if not hasattr(wrapper, "lora_A") or not hasattr(wrapper, "lora_B"): + return None + + if hasattr(wrapper, "disable_adapters") and wrapper.disable_adapters: + return None + if hasattr(wrapper, "merged") and wrapper.merged: + return None + + if not wrapper.lora_A: + return None + + if adapter_name not in wrapper.lora_A: + adapter_name = list(wrapper.lora_A.keys())[0] + + lora_A_module = wrapper.lora_A[adapter_name] + lora_B_module = wrapper.lora_B[adapter_name] + + weight_A = lora_A_module.weight # (E*R, dim1) + weight_B = lora_B_module.weight # (dim2, E*R) + scaling = wrapper.scaling[adapter_name] + num_experts = getattr(wrapper, "num_experts", 1) + + if experts_module is None: + experts_module = wrapper.get_base_layer() if hasattr(wrapper, "get_base_layer") else None + + # Model-specific LoRA extractor attached to the experts module, if any. + extractor_fn = getattr(experts_module, "_unsloth_lora_extractor_fn", None) + + if extractor_fn is not None: + return extractor_fn(wrapper, weight_A, weight_B, scaling, num_experts) + + return extract_moe_lora_weights_for_grouped_mm( + wrapper, + weight_A, + weight_B, + scaling, + num_experts, + experts_module=experts_module, + model_name="MoE", + ) + except Exception: + return None + + +def _extract_lora_weights( + param, adapter_name: str = "default", num_experts: int = None, experts_module=None +) -> Optional[Tuple[torch.Tensor, torch.Tensor, float]]: + """Compat wrapper around _extract_lora_from_wrapper; returns (first, second, scaling).""" + # This Unsloth Zoo code section is licensed under AGPL3 + + # Pass num_experts through so _extract_lora_from_wrapper can use it. + if num_experts is not None and not hasattr(param, "num_experts"): + param.num_experts = num_experts + + result = _extract_lora_from_wrapper(param, adapter_name, experts_module=experts_module) + if result is None: + return None + return result[0], result[1], result[2] + + +def _get_base_weight(param): + """Get base weight from a potentially wrapped parameter or module.""" + # This Unsloth Zoo code section is licensed under AGPL3 + + while hasattr(param, "base_layer"): + param = param.base_layer + + if HAS_BNB and isinstance(param, Params4bit): + if getattr(param, "quant_state", None) is None: + raise RuntimeError( + "unsloth: _get_base_weight saw a Params4bit with quant_state=None. " + "This usually means the model was used in forward before loading " + "completed quantization (meta placeholder still in place), or the " + "MoE quantizer patch did not fire for this expert. " + f"data.shape={tuple(param.data.shape)}, device={param.device}." + ) + return bnb.functional.dequantize_4bit(param.data, param.quant_state) + + if hasattr(param, "get_param"): + return param.get_param() + + if hasattr(param, "weight"): + return param.weight + + return param + + +def _get_lora_wrapper_for_param(experts_module, param_name): + """Get the PEFT ParamWrapper for gate_up_proj or down_proj; does not lazily set up wrappers.""" + # This Unsloth Zoo code section is licensed under AGPL3 + + if hasattr(experts_module, f"{param_name}_lora_wrapper"): + return getattr(experts_module, f"{param_name}_lora_wrapper") + + if hasattr(experts_module, param_name): + attr = getattr(experts_module, param_name) + if hasattr(attr, "lora_A"): # ParamWrapper + return attr + + return None + + +def native_moe_grouped_mm( + inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor +) -> torch.Tensor: + """Grouped_mm with backward fix for PyTorch's grouped_mm backward stride bug.""" + return _grouped_mm_with_backward_fix(inputs, weight, offsets) + + +def _apply_lora_grouped_mm( + inputs: torch.Tensor, + lora_B: torch.Tensor, + lora_A: torch.Tensor, + offsets: torch.Tensor, + scaling: float, + grouped_mm_func=native_moe_grouped_mm, +) -> torch.Tensor: + """Apply LoRA via grouped GEMM: result = ((X @ B) @ A) * scaling. + + inputs (total_tokens, in_dim); lora_B (E, in_dim, R); lora_A (E, R, out_dim). + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + # X @ B then result @ A; both already in native (E, ...) layout, no transpose. + lora_intermediate = grouped_mm_func(inputs, lora_B.contiguous(), offsets) + lora_delta = grouped_mm_func(lora_intermediate, lora_A.contiguous(), offsets) + + return lora_delta * scaling + + +def _should_use_separated_lora() -> bool: + """Use separated LoRA (default True); UNSLOTH_MOE_LORA_MERGED=1 forces the merged path.""" + return os.environ.get("UNSLOTH_MOE_LORA_MERGED", "0") != "1" + + +# Model-specific weight preprocessing hooks: each model registers a transposition +# function so the generic backend works across weight layouts. + +_WEIGHT_PREPROCESSORS = {} + + +def register_weight_preprocessor(model_type: str, preprocessor_fn): + """Register a weight preprocessor (weight, proj_type, hidden_dim) -> weight for a model type.""" + _WEIGHT_PREPROCESSORS[model_type] = preprocessor_fn + + +def get_weight_preprocessor(model_type: str): + """Get registered weight preprocessor for model type.""" + return _WEIGHT_PREPROCESSORS.get(model_type) + + +def preprocess_weight( + weight: torch.Tensor, proj_type: str, hidden_dim: int, model_type=None +): + """Preprocess a weight into (E, in_dim, out_dim) for grouped_mm. + + Uses a registered model-specific preprocessor if present, else transposes + by shape. proj_type is "gate_up" or "down". + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + if model_type and model_type in _WEIGHT_PREPROCESSORS: + return _WEIGHT_PREPROCESSORS[model_type](weight, proj_type, hidden_dim) + + if proj_type == "gate_up": + # Want (E, hidden_dim, 2*intermediate). + if weight.shape[1] == hidden_dim: + return weight + else: + return weight.transpose(-2, -1) + else: # down + # Want (E, intermediate, hidden_dim). + if weight.shape[2] == hidden_dim: + return weight + else: + return weight.transpose(-2, -1) + + +# Generic MoE detection and ParamWrapper patching. + + +def _normalize_model_type(value) -> str: + if value is None: + return "" + return str(value).lower().replace("-", "_") + + +def _iter_model_configs(model): + seen = set() + queue = [model] + while queue and len(seen) < 8: + current = queue.pop(0) + if current is None: + continue + current_id = id(current) + if current_id in seen: + continue + seen.add(current_id) + + config = getattr(current, "config", None) + if config is not None: + yield config + + for attr in ("base_model", "model"): + nested = getattr(current, attr, None) + if nested is not None and nested is not current: + queue.append(nested) + + +def _is_gpt_oss_model(model) -> bool: + for config in _iter_model_configs(model): + model_type = _normalize_model_type(getattr(config, "model_type", None)) + if model_type == "gpt_oss": + return True + + for attr in ("_name_or_path", "name_or_path"): + name = getattr(config, attr, None) + if name is None: + continue + # Match only the final path component so parent directories like + # /data/gpt-oss-tests/qwen3-7b do not count as gpt-oss. + base = str(name).replace("\\", "/").rstrip("/").rsplit("/", 1)[-1] + if "gpt_oss" in _normalize_model_type(base): + return True + + return False + + +def _set_gpt_oss_grouped_mm_format_on_experts(module) -> bool: + if module is None: + return False + if module.__class__.__name__ != "GptOssExperts": + return False + if bool(getattr(module, "_unsloth_grouped_mm_format", False)): + return False + # Require the gpt-oss (E, in, out) weight signature: gate_up's out dim is + # twice down's in dim. Same-named classes with other layouts stay unflagged. + gate_shape = _get_param_shape_from_module(module, "gate_up_proj") + down_shape = _get_param_shape_from_module(module, "down_proj") + if gate_shape is None or down_shape is None: + return False + if len(gate_shape) < 3 or len(down_shape) < 3: + return False + if gate_shape[0] != down_shape[0]: + return False + if gate_shape[-2] != down_shape[-1] or gate_shape[-1] != 2 * down_shape[-2]: + return False + module._unsloth_grouped_mm_format = True + return True + + +def patch_gpt_oss_grouped_mm_format(model) -> int: + """ + Mark GPT-OSS experts as storing weights in grouped_mm format. + + Stock transformers GPT-OSS experts use (E, in_dim, out_dim) tensors but do + not carry Unsloth's `_unsloth_grouped_mm_format` instance flag. Set it on + live expert modules so the shared MoE LoRA extractor chooses GPT-OSS + ordering instead of the Qwen-style fallback. + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + if model is None or not _is_gpt_oss_model(model): + return 0 + + modules = getattr(model, "modules", None) + if not callable(modules): + return 0 + + updated = 0 + for module in modules(): + if _set_gpt_oss_grouped_mm_format_on_experts(module): + updated += 1 + return updated + + +def _patch_peft_get_peft_model_for_moe(): + # This Unsloth Zoo code section is licensed under AGPL3 + + global _original_peft_get_peft_model + if _original_peft_get_peft_model is not None: + return + + try: + import peft + except Exception: + return + + original_get_peft_model = getattr(peft, "get_peft_model", None) + if original_get_peft_model is None: + return + if getattr(original_get_peft_model, "_unsloth_moe_patched", False): + return + + _original_peft_get_peft_model = original_get_peft_model + + @wraps(original_get_peft_model) + def patched_get_peft_model(model, *args, **kwargs): + peft_model = original_get_peft_model(model, *args, **kwargs) + try: + patch_gpt_oss_grouped_mm_format(model) + if peft_model is not model: + patch_gpt_oss_grouped_mm_format(peft_model) + except Exception: + pass + return peft_model + + patched_get_peft_model._unsloth_moe_patched = True + peft.get_peft_model = patched_get_peft_model + + for module_name in ("peft.mapping_func", "peft.mapping"): + try: + module = importlib.import_module(module_name) + except Exception: + continue + if getattr(module, "get_peft_model", None) is original_get_peft_model: + module.get_peft_model = patched_get_peft_model + + +def _is_moe_experts_module(module) -> bool: + """Generic check for an MoE experts layer with stacked 3D expert weights. + + Matches gate_up_proj/down_proj (Qwen3-MoE etc.) or w1/w2/w3 (older models). + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + import torch.nn as nn + + # After PEFT's parametrize wrapping, gate_up_proj is a Tensor (not Parameter), + # so accept both. + if hasattr(module, "gate_up_proj"): + param = module.gate_up_proj + # 4-bit params are packed into 2D tensors. + if HAS_BNB and isinstance(param, Params4bit) and param.ndim == 2: + return True + # Standard MoE weights are 3D (num_experts, in, out). + if isinstance(param, (nn.Parameter, torch.Tensor)) and param.ndim in (2, 3): + return True + + # w1/w2 pattern (separate gate/up projections). + if hasattr(module, "w1") and hasattr(module, "w2"): + w1 = module.w1 + if isinstance(w1, (nn.Parameter, torch.Tensor)) and w1.ndim in (2, 3): + return True + + return False + + +# Aliases for compatibility with gpt_oss.py +_get_moe_lora_weights = _extract_lora_from_wrapper + + +# Store original ParamWrapper.forward for fallback +_original_param_wrapper_forward = None + + +def _patched_param_wrapper_forward( + self, x: torch.Tensor, *args, **kwargs +) -> torch.Tensor: + """Patched ParamWrapper.forward for MoE separated LoRA. + + For MoE experts: bypass PEFT's _activate_lora and stash LoRA data by + parameter_name for forward_native_grouped_mm. For non-MoE: original forward. + """ + # This Unsloth Zoo code section is licensed under AGPL3 + + # Use self.base_layer (immediate parent), NOT get_base_layer() which recurses + # to the deepest layer; the wrapper chain down_proj -> gate_up_proj -> + # Qwen3MoeExperts must be preserved. + immediate_base_layer = self.base_layer + + # For stashing LoRA data we need the actual experts module (recursive lookup). + experts_module = self.get_base_layer() + + use_separated = _should_use_separated_lora() + param_name = getattr(self, "parameter_name", None) + + if ( + use_separated + and param_name in ("gate_up_proj", "down_proj") + and _is_moe_experts_module(experts_module) + ): + # MoE experts: bypass PEFT's _activate_lora, use separated computation. + if self.disable_adapters: + if self.merged: + self.unmerge() + return immediate_base_layer(x, *args, **kwargs) + + if self.merged: + return immediate_base_layer(x, *args, **kwargs) + + # Ensure wrapper.num_experts is set for LoRA weight reshaping. + if not hasattr(self, "num_experts"): + if hasattr(experts_module, "num_experts"): + self.num_experts = experts_module.num_experts + elif hasattr(experts_module, param_name): + p = getattr(experts_module, param_name) + if hasattr(p, "shape") and len(p.shape) >= 1: + self.num_experts = p.shape[0] + + # Extract LoRA for this parameter and stash on the experts module + # (not base_layer): _unsloth_lora_gate_up_proj / _unsloth_lora_down_proj. + lora_data = _extract_lora_from_wrapper(self) + + if lora_data is not None and param_name: + lora_attr = f"_unsloth_lora_{param_name}" + setattr(experts_module, lora_attr, lora_data) + + try: + # Immediate base_layer preserves the wrapper chain. + result = immediate_base_layer(x, *args, **kwargs) + finally: + if param_name: + lora_attr = f"_unsloth_lora_{param_name}" + if hasattr(experts_module, lora_attr): + delattr(experts_module, lora_attr) + + return result + + # Non-MoE: original PEFT forward with _activate_lora. + return _original_param_wrapper_forward(self, x, *args, **kwargs) + + +def patch_param_wrapper_for_moe(): + """Patch PEFT's ParamWrapper.forward for MoE separated LoRA (call after PEFT import).""" + # This Unsloth Zoo code section is licensed under AGPL3 + + global _original_param_wrapper_forward + + module = _load_cached_moe_utils_module() + if module is not None and hasattr(module, "patch_param_wrapper_for_moe"): + try: + return module.patch_param_wrapper_for_moe() + except Exception: + pass + + try: + from peft.tuners.lora.layer import ParamWrapper + + if _original_param_wrapper_forward is None: + _original_param_wrapper_forward = ParamWrapper.forward + + ParamWrapper.forward = _patched_param_wrapper_forward + _patch_peft_get_peft_model_for_moe() + + return True + except ImportError: + return False + + +def forward_native_grouped_mm( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Native PyTorch grouped-GEMM MoE forward via torch._grouped_mm (no Triton; needs runtime support).""" + # This Unsloth Zoo code section is licensed under AGPL3 + + # Runtime safety check (defense in depth). + if not _check_torch_grouped_mm_supported(): + major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) + raise RuntimeError( + f"torch._grouped_mm is not supported on this device (Compute Capability {major}.{minor}). " + f"Set UNSLOTH_MOE_BACKEND='unsloth_triton' or 'native_torch' to use a compatible backend." + ) + + is_2d_input = hidden_states.dim() == 2 + if is_2d_input: + sequence_length, hidden_dim = hidden_states.shape + batch_size = 1 + else: + batch_size, sequence_length, hidden_dim = hidden_states.shape + + hidden_states = hidden_states.view(-1, hidden_dim) + + # Routing: count tokens per expert, sort to group by expert, gather inputs. + flat_top_k = top_k_index.view(-1) + num_tokens_per_expert = torch.bincount(flat_top_k, minlength=self.num_experts).int() + sorted_indices = torch.argsort(flat_top_k, stable=True) + token_indices = sorted_indices // top_k_index.shape[-1] + permuted_input = hidden_states[token_indices] + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) + + # Gate + Up projection with optional separated LoRA (default). + use_separated_lora = _should_use_separated_lora() + gate_up_lora = None + + # Prefer LoRA injected by the patched ParamWrapper; fall back to the parameter. + if getattr(self, "_unsloth_lora_gate_up_proj", None) is not None: + gate_up_lora = self._unsloth_lora_gate_up_proj[:3] # (first, second, scaling) + elif ( + use_separated_lora + and hasattr(self, "gate_up_proj") + and _has_lora_adapters(self.gate_up_proj) + ): + gate_up_lora = _extract_lora_weights( + self.gate_up_proj, num_experts=self.num_experts, experts_module=self + ) + + if hasattr(self, "gate_up_proj"): + gate_up_base = _get_base_weight(self.gate_up_proj) + model_type = getattr(self, "_unsloth_model_type", None) + + # grouped_mm backward needs contiguous weights; preprocess may return a transposed view. + w1 = preprocess_weight(gate_up_base, "gate_up", hidden_dim, model_type) + mm1_out = _grouped_mm_with_backward_fix(permuted_input, w1, offsets) + + # Separated LoRA: + ((X @ first) @ second) * scaling. + if gate_up_lora is not None: + first_weight, second_weight, scaling = gate_up_lora + + # Cast to input dtype (LoRA is float32) and make contiguous for grouped_mm. + first_weight = first_weight.to(permuted_input.dtype).contiguous() + second_weight = second_weight.to(permuted_input.dtype).contiguous() + + try: + lora_out = _grouped_mm_with_backward_fix(permuted_input, first_weight, offsets) + lora_out = lora_out.contiguous() + except RuntimeError as e: + raise e + + # Second matmul; pad an unaligned output dim or fall back on failure. + try: + if second_weight.shape[-1] % 8 != 0: + pad_size = 8 - (second_weight.shape[-1] % 8) + second_weight_padded = F.pad( + second_weight, (0, pad_size) + ).contiguous() + lora_delta = _grouped_mm_with_backward_fix( + lora_out, second_weight_padded, offsets + ) + lora_delta = lora_delta[:, :-pad_size] + else: + lora_delta = _grouped_mm_with_backward_fix( + lora_out, second_weight, offsets + ) + except RuntimeError: + # Manual loop fallback on grouped_mm failure (e.g. stride alignment). + lora_delta = torch.empty( + (lora_out.shape[0], second_weight.shape[-1]), + dtype=lora_out.dtype, + device=lora_out.device, + ) + cpu_offsets = offsets.cpu().tolist() + prev_offset = 0 + for i, end in enumerate(cpu_offsets): + if prev_offset < end: + lora_delta[prev_offset:end] = torch.matmul( + lora_out[prev_offset:end], second_weight[i] + ) + prev_offset = end + + mm1_out = mm1_out + lora_delta * scaling + + if hasattr(self, "gate_up_proj_bias") and self.gate_up_proj_bias is not None: + num_repeats = num_tokens_per_expert.to(self.gate_up_proj_bias.device) + bias_expanded = self.gate_up_proj_bias.repeat_interleave(num_repeats, dim=0) + mm1_out = mm1_out + bias_expanded.to(mm1_out.dtype) + + if "GptOssExperts" in self.__class__.__name__: + gate = mm1_out[..., ::2] + up = mm1_out[..., 1::2] + else: + gate, up = mm1_out.chunk(2, dim=-1) + + elif hasattr(self, "w1") and hasattr(self, "w3"): + # Separate w1/w3 weights (older models). + w1_base = _get_base_weight(self.w1) + w3_base = _get_base_weight(self.w3) + + w1 = w1_base.transpose(-2, -1) + w3 = w3_base.transpose(-2, -1) + + gate = _grouped_mm_with_backward_fix(permuted_input, w1, offsets) + up = _grouped_mm_with_backward_fix(permuted_input, w3, offsets) + + # Add LoRA for w1 and w3 separately if present. + if use_separated_lora: + if _has_lora_adapters(self.w1): + w1_lora = _extract_lora_weights(self.w1, experts_module=self) + if w1_lora is not None: + lora_A, lora_B, scaling = w1_lora + lora_A_t = lora_A.transpose(-2, -1) + lora_A_out = _grouped_mm_with_backward_fix( + permuted_input, lora_A_t, offsets + ) + lora_B_t = lora_B.transpose(-2, -1) + lora_B_out = _grouped_mm_with_backward_fix(lora_A_out, lora_B_t, offsets) + gate = gate + lora_B_out * scaling + + if _has_lora_adapters(self.w3): + w3_lora = _extract_lora_weights(self.w3, experts_module=self) + if w3_lora is not None: + lora_A, lora_B, scaling = w3_lora + lora_A_t = lora_A.transpose(-2, -1) + lora_A_out = _grouped_mm_with_backward_fix( + permuted_input, lora_A_t, offsets + ) + lora_B_t = lora_B.transpose(-2, -1) + lora_B_out = _grouped_mm_with_backward_fix(lora_A_out, lora_B_t, offsets) + up = up + lora_B_out * scaling + else: + raise AttributeError("MoE layer must have 'gate_up_proj' or 'w1'/'w3'.") + + # Activation + if "GptOssExperts" in self.__class__.__name__: + # Custom GptOss activation. + limit = getattr(self, "limit", 7.0) + alpha = getattr(self, "alpha", 1.702) + + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + inter = (up + 1.0) * glu + elif hasattr(self, 'act_fn') and callable(self.act_fn): + inter = self.act_fn(gate) * up + else: + inter = F.silu(gate) * up + + # Down projection with optional separated LoRA (default). + down_lora = None + + # Prefer LoRA injected by the patched ParamWrapper; fall back to the parameter. + if getattr(self, "_unsloth_lora_down_proj", None) is not None: + down_lora = self._unsloth_lora_down_proj[:3] # (first, second, scaling) + elif ( + use_separated_lora + and hasattr(self, "down_proj") + and _has_lora_adapters(self.down_proj) + ): + down_lora = _extract_lora_weights(self.down_proj, num_experts=self.num_experts, experts_module=self) + + if hasattr(self, "down_proj"): + down_base = _get_base_weight(self.down_proj) + model_type = getattr(self, "_unsloth_model_type", None) + w2 = preprocess_weight(down_base, "down", hidden_dim, model_type) + mm2_out = _grouped_mm_with_backward_fix(inter, w2, offsets) + + if down_lora is not None: + first_weight, second_weight, scaling = down_lora + + # Cast to input dtype (LoRA is float32) and make contiguous for grouped_mm. + first_weight = first_weight.to(inter.dtype).contiguous() + second_weight = second_weight.to(inter.dtype).contiguous() + + lora_out = _grouped_mm_with_backward_fix(inter, first_weight, offsets) + lora_out = lora_out.contiguous() + + try: + lora_delta = _grouped_mm_with_backward_fix(lora_out, second_weight, offsets) + except RuntimeError: + # Manual loop fallback. + lora_delta = torch.empty( + (lora_out.shape[0], second_weight.shape[-1]), + dtype=lora_out.dtype, + device=lora_out.device, + ) + cpu_offsets = offsets.cpu().tolist() + prev_offset = 0 + for i, end in enumerate(cpu_offsets): + if prev_offset < end: + lora_delta[prev_offset:end] = torch.matmul( + lora_out[prev_offset:end], second_weight[i] + ) + prev_offset = end + + mm2_out = mm2_out + lora_delta * scaling + + if hasattr(self, "down_proj_bias") and self.down_proj_bias is not None: + bias_expanded = self.down_proj_bias.repeat_interleave( + num_tokens_per_expert.to(self.down_proj_bias.device), dim=0 + ).to(mm2_out.device) + mm2_out = mm2_out + bias_expanded.to(mm2_out.dtype) + + elif hasattr(self, "w2"): + w2_base = _get_base_weight(self.w2) + w2 = w2_base.transpose(-2, -1) + mm2_out = _grouped_mm_with_backward_fix(inter, w2, offsets) + + if use_separated_lora and _has_lora_adapters(self.w2): + w2_lora = _extract_lora_weights(self.w2, experts_module=self) + if w2_lora is not None: + lora_A, lora_B, scaling = w2_lora + lora_A_t = lora_A.transpose(-2, -1).contiguous() + lora_A_out = _grouped_mm_with_backward_fix(inter, lora_A_t, offsets) + lora_B_t = lora_B.transpose(-2, -1).contiguous() + lora_B_out = _grouped_mm_with_backward_fix(lora_A_out, lora_B_t, offsets) + mm2_out = mm2_out + lora_B_out * scaling + else: + raise AttributeError("MoE layer must have 'down_proj' or 'w2'.") + + # Apply routing weights and scatter-add (reduce). + flat_weights = top_k_weights.view(-1) + permuted_weights = flat_weights[sorted_indices] + mm2_out = mm2_out * permuted_weights.unsqueeze(-1) + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + final_hidden_states.index_add_(0, token_indices, mm2_out.to(hidden_states.dtype)) + + if is_2d_input: + return final_hidden_states + + return final_hidden_states.view(batch_size, sequence_length, hidden_dim) + + +def forward_triton_grouped_gemm( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Grouped-GEMM MoE forward via Triton kernels (torch.compile-compatible, mode="max-autotune").""" + # This Unsloth Zoo code section is licensed under AGPL3 + + from unsloth.kernels.moe.grouped_gemm.interface import grouped_gemm + from unsloth.kernels.moe.autotune_cache import get_or_autotune_moe_kernels + + if not hasattr(self, "_unsloth_moe_configs"): + self._unsloth_moe_configs = None + + use_separated_lora = _should_use_separated_lora() + + # gate_up LoRA from the patched ParamWrapper (mirrors the down block below). + gate_up_lora = None + if getattr(self, "_unsloth_lora_gate_up_proj", None) is not None: + gate_up_lora = self._unsloth_lora_gate_up_proj[:3] + elif ( + use_separated_lora + and hasattr(self, "gate_up_proj") + and _has_lora_adapters(self.gate_up_proj) + ): + gate_up_lora = _extract_lora_weights( + self.gate_up_proj, num_experts=self.num_experts + ) + + # Flatten 3D inputs (batch_size, seq_len, hidden_dim). + is_3d = hidden_states.dim() == 3 + if is_3d: + batch_size, seq_len, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + num_tokens = batch_size * seq_len + if top_k_index.dim() == 3: + top_k_index = top_k_index.view(-1, top_k_index.shape[-1]) + if top_k_weights.dim() == 3: + top_k_weights = top_k_weights.view(-1, top_k_weights.shape[-1]) + else: + num_tokens, hidden_dim = hidden_states.shape + + top_k = top_k_index.shape[1] + + # Cache model dims and kernel configs on first call. + if self._unsloth_moe_configs is None: + intermediate_dim = self.gate_up_proj.shape[1] // 2 + + # Autotune first GEMM. + gemm1_configs = get_or_autotune_moe_kernels( + num_experts=self.num_experts, + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim * 2, + top_k=top_k, + dtype=hidden_states.dtype, + ) + + # Autotune second GEMM (output dim is hidden_dim). + gemm2_configs = get_or_autotune_moe_kernels( + num_experts=self.num_experts, + hidden_dim=intermediate_dim, + intermediate_dim=hidden_dim, + top_k=top_k, + dtype=hidden_states.dtype, + ) + + self._unsloth_moe_configs = (intermediate_dim, gemm1_configs, gemm2_configs) + torch.cuda.empty_cache() + + intermediate_dim, gemm1_configs, gemm2_configs = self._unsloth_moe_configs + fwd_config_1, bwd_dX_config_1, bwd_dW_config_1 = gemm1_configs + fwd_config_2, bwd_dX_config_2, bwd_dW_config_2 = gemm2_configs + + token_counts_by_expert, gather_indices = _get_routing_indices( + top_k_index, self.num_experts + ) + offsets = torch.cumsum(token_counts_by_expert, dim=0, dtype=torch.int32) + + if self.gate_up_proj.shape[-1] == hidden_dim: + w1 = self.gate_up_proj + else: + w1 = self.gate_up_proj.transpose(-2, -1).contiguous() + + # First grouped GEMM: gate_up projection. + first_gemm_output = grouped_gemm( + X=hidden_states, + W=w1, + m_sizes=token_counts_by_expert, + topk=top_k, + gather_indices=gather_indices, + permute_x=True, + permute_y=False, + autotune=False, # cached configs + kernel_config_fwd=fwd_config_1, + kernel_config_bwd_dX=bwd_dX_config_1, + kernel_config_bwd_dW=bwd_dW_config_1, + is_first_gemm=True, + ) + + # Separated LoRA for gate_up. grouped_gemm ran permute_x=True so first_gemm_output + # is expert-sorted; _apply_lora_grouped_mm wants pre-permuted input, so gather via + # gather_indices // top_k (expert-sorted row -> originating token row). + if gate_up_lora is not None: + first_weight, second_weight, scaling = gate_up_lora + first_weight = first_weight.to(hidden_states.dtype) + second_weight = second_weight.to(hidden_states.dtype) + permuted_hidden = hidden_states[gather_indices // top_k] + gate_up_lora_delta = _apply_lora_grouped_mm( + permuted_hidden, + first_weight, + second_weight, + offsets, + scaling, + grouped_mm_func=native_moe_grouped_mm, + ) + first_gemm_output = first_gemm_output + gate_up_lora_delta + + # Activation + gate*up. + if hasattr(self, 'act_fn') and callable(self.act_fn): + gate, up = first_gemm_output.chunk(2, dim=-1) + intermediate = self.act_fn(gate) * up + else: + intermediate = _silu_and_mul(first_gemm_output) + + # Grouped GEMM 2: down projection. + down_lora = None + if getattr(self, "_unsloth_lora_down_proj", None) is not None: + down_lora = self._unsloth_lora_down_proj[:3] + elif ( + use_separated_lora + and hasattr(self, "down_proj") + and _has_lora_adapters(self.down_proj) + ): + down_lora = _extract_lora_weights(self.down_proj, num_experts=self.num_experts) + + if self.down_proj.shape[-1] == intermediate.shape[-1]: + w2 = self.down_proj + else: + w2 = self.down_proj.transpose(-2, -1).contiguous() + + second_gemm_output = grouped_gemm( + X=intermediate, + W=w2, + m_sizes=token_counts_by_expert, + topk=top_k, + gather_indices=gather_indices, + permute_x=False, + permute_y=True, + autotune=False, # cached configs + kernel_config_fwd=fwd_config_2, + kernel_config_bwd_dX=bwd_dX_config_2, + kernel_config_bwd_dW=bwd_dW_config_2, + is_first_gemm=False, + ) + + # Separated LoRA for down (intermediate already permuted from step 1, same offsets). + if down_lora is not None: + first_weight, second_weight, scaling = down_lora + + first_weight = first_weight.to(intermediate.dtype) + second_weight = second_weight.to(intermediate.dtype) + + lora_delta = _apply_lora_grouped_mm( + intermediate, + first_weight, + second_weight, + offsets, + scaling, + grouped_mm_func=native_moe_grouped_mm + ) + + second_gemm_output = second_gemm_output + lora_delta + + # Apply routing weights and sum across top_k: (num_tokens, top_k, hidden) -> (num_tokens, hidden). + top_k_weights_casted = top_k_weights.to(hidden_states.dtype) + final_hidden_states = ( + second_gemm_output.view(num_tokens, top_k, hidden_dim) + * top_k_weights_casted[..., None] + ) + final_hidden_states = final_hidden_states.sum(dim=1) + + if is_3d: + final_hidden_states = final_hidden_states.view(batch_size, seq_len, hidden_dim) + + return final_hidden_states + + +@torch.compiler.disable +def forward_native_moe_loop( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Loop over experts with routed tokens; torch.compile-disabled to avoid graph breaks on dynamic control flow.""" + # This Unsloth Zoo code section is licensed under AGPL3 + final_hidden_states = torch.zeros_like(hidden_states) + use_separated_lora = _should_use_separated_lora() + + gate_up_lora = getattr(self, "_unsloth_lora_gate_up_proj", None) + if gate_up_lora is not None: + gate_up_lora = gate_up_lora[:3] + elif ( + use_separated_lora + and hasattr(self, "gate_up_proj") + and _has_lora_adapters(self.gate_up_proj) + ): + gate_up_lora = _extract_lora_weights( + self.gate_up_proj, num_experts=self.num_experts, experts_module=self + ) + # Pre-cast LoRA factors to the activation dtype once (avoid per-expert .to()). + # `scaling` is left alone: a Python float is a no-op, a tensor broadcasts. + if gate_up_lora is not None: + _gate_up_first, _gate_up_second, _gate_up_scaling = gate_up_lora + gate_up_lora = ( + _gate_up_first.to(hidden_states.dtype), + _gate_up_second.to(hidden_states.dtype), + _gate_up_scaling, + ) + + down_lora = getattr(self, "_unsloth_lora_down_proj", None) + if down_lora is not None: + down_lora = down_lora[:3] + elif ( + use_separated_lora + and hasattr(self, "down_proj") + and _has_lora_adapters(self.down_proj) + ): + down_lora = _extract_lora_weights( + self.down_proj, num_experts=self.num_experts, experts_module=self + ) + if down_lora is not None: + _down_first, _down_second, _down_scaling = down_lora + down_lora = ( + _down_first.to(hidden_states.dtype), + _down_second.to(hidden_states.dtype), + _down_scaling, + ) + + # Expert mask -> which experts have tokens. + with torch.no_grad(): + expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) # (num_experts, top_k, n_tokens) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + # Some patches (Qwen3-VL-MoE) store experts in grouped_mm layout (E, in, out) + # rather than F.linear's (E, out, in) and set _unsloth_grouped_mm_format=True. + # Prefer it over the shape check, which is unsafe when intermediate_dim == hidden_dim. + grouped_mm_format = bool(getattr(self, "_unsloth_grouped_mm_format", False)) + + # GPT-OSS uses interleaved gate/up, clamped swiglu, and per-expert biases. + is_gpt_oss = "GptOssExperts" in self.__class__.__name__ + + for expert_idx_t in expert_hit: + expert_idx = expert_idx_t.item() + + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + + # gate_up projection for this expert ('gate_up_proj' or 'w1'/'w3'). + if hasattr(self, "gate_up_proj"): + gate_up_weight = self.gate_up_proj[expert_idx] + if grouped_mm_format or gate_up_weight.shape[-1] != current_state.shape[-1]: + gate_up_weight = gate_up_weight.T + gate_up = F.linear(current_state, gate_up_weight) + if gate_up_lora is not None: + first_weight, second_weight, scaling = gate_up_lora + lora_delta = current_state @ first_weight[expert_idx] + lora_delta = lora_delta @ second_weight[expert_idx] + gate_up = gate_up + lora_delta * scaling + if is_gpt_oss: + gate_up_bias = getattr(self, "gate_up_proj_bias", None) + if gate_up_bias is not None: + gate_up = gate_up + gate_up_bias[expert_idx].to(gate_up.dtype) + gate = gate_up[..., ::2] + up = gate_up[..., 1::2] + else: + gate, up = gate_up.chunk(2, dim=-1) + else: + gate = F.linear(current_state, self.w1[expert_idx]) + up = F.linear(current_state, self.w3[expert_idx]) + + if is_gpt_oss: + limit = getattr(self, "limit", 7.0) + alpha = getattr(self, "alpha", 1.702) + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=-limit, max=limit) + current_hidden_states = (up + 1.0) * (gate * torch.sigmoid(gate * alpha)) + elif hasattr(self, "act_fn") and callable(self.act_fn): + current_hidden_states = self.act_fn(gate) * up + else: + current_hidden_states = F.silu(gate) * up + + # down projection for this expert. + if hasattr(self, "down_proj"): + down_weight = self.down_proj[expert_idx] + # Mirror gate_up: prefer the flag over the shape heuristic (unsafe at square dims). + if grouped_mm_format or down_weight.shape[-1] != current_hidden_states.shape[-1]: + down_weight = down_weight.T + down = F.linear(current_hidden_states, down_weight) + if down_lora is not None: + first_weight, second_weight, scaling = down_lora + lora_delta = current_hidden_states @ first_weight[expert_idx] + lora_delta = lora_delta @ second_weight[expert_idx] + down = down + lora_delta * scaling + if is_gpt_oss: + down_bias = getattr(self, "down_proj_bias", None) + if down_bias is not None: + down = down + down_bias[expert_idx].to(down.dtype) + current_hidden_states = down + else: + current_hidden_states = F.linear(current_hidden_states, self.w2[expert_idx]) + + current_hidden_states = ( + current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + ) + + final_hidden_states.index_add_( + 0, token_idx, current_hidden_states.to(final_hidden_states.dtype) + ) + + return final_hidden_states diff --git a/unsloth_compiled_cache/unsloth_compiled_module_gemma4.py b/unsloth_compiled_cache/unsloth_compiled_module_gemma4.py new file mode 100644 index 0000000000000000000000000000000000000000..fa8d4c309ed7c7578a288f5b4dbabfccd631dd80 --- /dev/null +++ b/unsloth_compiled_cache/unsloth_compiled_module_gemma4.py @@ -0,0 +1,1706 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + + +from unsloth_zoo.loss_utils import ( + fused_linear_cross_entropy, + unsloth_fused_ce_loss, +) + +scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention +@torch.compiler.disable(recursive = False) +def disable_compile_scaled_dot_product_attention(*args, **kwargs): + return scaled_dot_product_attention(*args, **kwargs) +pass + + +from transformers.modeling_flash_attention_utils import is_flash_attn_available + +if is_flash_attn_available(): + try: + from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask + except: + flash_attn_supports_top_left_mask = None + try: + from transformers.modeling_flash_attention_utils import _flash_attention_forward + except: + _flash_attention_forward = None + try: + from transformers.modeling_flash_attention_utils import FlashAttentionKwargs + except: + FlashAttentionKwargs = None + try: + from transformers.modeling_flash_attention_utils import flash_attn_varlen_func + except: + flash_attn_varlen_func = None +else: + flash_attn_supports_top_left_mask = None + _flash_attention_forward = None + FlashAttentionKwargs = None + flash_attn_varlen_func = None +pass + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} + +from torch.nn import CrossEntropyLoss + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def normal_cross_entropy_loss(self, hidden_states, labels): + logits = self.lm_head(hidden_states) + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + return loss, logits +pass + +# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie +# os.environ['UNSLOTH_RETURN_LOGITS'] = '1' +LOGITS_ERROR_STRING = \ + "Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\ + 'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n'\ + "```\nimport os\n"\ + "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\ + "trainer.train()\n```\n"\ + "No need to restart your console - just add `os.environ['UNSLOTH_RETURN_LOGITS'] = '1'` before trainer.train() and re-run the cell!" + +def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING) +def return_none(*args, **kwargs): return None +class EmptyLogits: + def __init__(self): return + def raise_getattr_error(self, attr): return return_none if attr == "to" else raise_logits_error + __getitem__ = raise_logits_error + __getattr__ = raise_getattr_error + def __repr__(self): return LOGITS_ERROR_STRING + def __str__ (self): return LOGITS_ERROR_STRING + # Stateless pickling so accelerate gather_object works on the sentinel + def __reduce__(self): return (type(self), ()) + # Gathered copies must compare equal in accelerate debug mode + def __eq__(self, other): return type(other).__name__ == "EmptyLogits" + __hash__ = object.__hash__ +pass +EMPTY_LOGITS = EmptyLogits() +functions = dir(torch.Tensor) +for j, function in enumerate(functions): + if function.startswith("__") and function.endswith("__"): + exec(f"def raise_{j}(*args, **kwargs): print('{function}')", globals(), locals()) + try: exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals()) + except: continue +pass +# The loop above stomps pickle hooks with stubs returning None; restore them. +for function in ("__reduce__", "__reduce_ex__", "__getstate__", "__setstate__"): + try: delattr(EMPTY_LOGITS, function) + except Exception: pass +pass + + +def mask_attention_mask_out(labels = None, attention_mask = None): + if labels is not None and attention_mask is not None: + attention_mask = attention_mask.to(device = labels.device) + labels[attention_mask == 0] = -100 + return labels +pass + + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.gemma4.modeling_gemma4 import (F, math, Callable, Optional, torch, nn, init, ACT2FN, Cache, PreTrainedConfig, GenerationMixin, create_causal_mask, create_sliding_window_causal_mask, FlashAttentionKwargs, BaseModelOutputWithPast, ModelOutput, CausalLMOutputWithPast, ROPE_INIT_FUNCTIONS, dynamic_rope_update, ALL_ATTENTION_FUNCTIONS, PreTrainedModel, Unpack, TransformersKwargs, can_return_tuple, maybe_autocast, Gemma4AudioConfig, Gemma4Config, Gemma4TextConfig, Gemma4VisionConfig, Gemma4CausalLMOutputWithPast, Gemma4AudioCausalConv1d, Gemma4PreTrainedModel, Gemma4TextModel, Gemma4ForCausalLM, Gemma4Model, Gemma4ForConditionalGeneration, __name__, Gemma4TextExperts, create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask) + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4ClippableLinear_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_clipped_linears: + hidden_states = torch.clamp(hidden_states, self.input_min, self.input_max) + + hidden_states = self.linear(hidden_states) + + if self.use_clipped_linears: + hidden_states = torch.clamp(hidden_states, self.output_min, self.output_max) + + return hidden_states + +class Gemma4ClippableLinear(nn.Module): + def __init__( + self, + config: Gemma4VisionConfig | Gemma4AudioConfig, + in_features: int, + out_features: int, + ) -> None: + super().__init__() + self.use_clipped_linears = config.use_clipped_linears + self.linear = nn.Linear(in_features, out_features, bias=False) + + if self.use_clipped_linears: + self.register_buffer("input_min", torch.tensor(-float("inf"))) + self.register_buffer("input_max", torch.tensor(float("inf"))) + self.register_buffer("output_min", torch.tensor(-float("inf"))) + self.register_buffer("output_max", torch.tensor(float("inf"))) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return Gemma4ClippableLinear_forward(self, hidden_states=hidden_states) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def Gemma4RMSNorm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + normed_output = self._norm(hidden_states.float()) + if self.with_scale: + normed_output = normed_output * self.weight.float() + return normed_output.type_as(hidden_states) + +class Gemma4RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True): + super().__init__() + self.eps = eps + self.with_scale = with_scale + + if self.with_scale: + self.weight = nn.Parameter(torch.ones(dim), requires_grad=True) + + def _norm(self, hidden_states: torch.Tensor): + mean_squared = hidden_states.pow(2).mean(-1, keepdim=True) + self.eps + # Use torch.pow() (over torch.sqrt() or torch.rsqrt()) to addess compiler differences between Torch and JAX + return hidden_states * torch.pow(mean_squared, -0.5) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return Gemma4RMSNorm_forward(self, hidden_states=hidden_states) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +@torch.no_grad() +def Gemma4AudioRelPositionalEncoding_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + position_ids = torch.arange(12, -1, -1, device=hidden_states.device) + position_ids = position_ids[..., None] + scaled_time = position_ids * self.inv_timescales.to(device=hidden_states.device) + pos_embed = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=-1) + return pos_embed.to(dtype=hidden_states.dtype) + +class Gemma4AudioRelPositionalEncoding(nn.Module): + """Sinusoidal relative positional encoding for the audio encoder. + + Produces position embeddings of shape [1, 2*context_size - 1, hidden_size] with + concatenated [sin..., cos...] layout matching the original Gemma4 convention. + """ + + inv_timescales: torch.Tensor + + def __init__(self, config: Gemma4AudioConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.context_size = ( + config.attention_chunk_size + config.attention_context_left - 1 + config.attention_context_right + ) + min_timescale = 1.0 + max_timescale = 10000.0 + num_timescales = self.hidden_size // 2 + log_timescale_increment = math.log(max_timescale / min_timescale) / max(num_timescales - 1, 1) + inv_timescales = min_timescale * torch.exp(torch.arange(num_timescales) * -log_timescale_increment) + self.register_buffer("inv_timescales", inv_timescales.unsqueeze(0).unsqueeze(0), persistent=False) + + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return Gemma4AudioRelPositionalEncoding_forward(self, hidden_states=hidden_states) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def Gemma4AudioSubSampleConvProjectionLayer_forward(self, hidden_states: torch.Tensor, mask: torch.Tensor | None = None): + if mask is not None: + mask = mask.to(device=hidden_states.device) + hidden_states = hidden_states * mask[:, None, :, None] + + hidden_states = self.conv(hidden_states.to(self.conv.weight.dtype)) + hidden_states = self.act(self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2).contiguous()) + + if mask is not None: + mask = mask[:, ::2] + + return hidden_states, mask + +class Gemma4AudioSubSampleConvProjectionLayer(nn.Module): + def __init__(self, in_channels, out_channels, norm_eps): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(3, 3), + stride=(2, 2), + padding=1, + bias=False, + ) + self.norm = nn.LayerNorm(out_channels, eps=norm_eps, elementwise_affine=True, bias=False) + self.act = nn.ReLU() + + def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor | None = None): + return Gemma4AudioSubSampleConvProjectionLayer_forward(self, hidden_states=hidden_states, mask=mask) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4AudioSubSampleConvProjection_forward( + self, + input_features: torch.Tensor, + input_features_mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states = input_features.unsqueeze(1) + hidden_states, mask = self.layer0(hidden_states, input_features_mask) + hidden_states, mask = self.layer1(hidden_states, mask) + + batch_size, _, seq_len, _ = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 1).contiguous().reshape(batch_size, seq_len, -1) + return self.input_proj_linear(hidden_states), mask + +class Gemma4AudioSubSampleConvProjection(nn.Module): + def __init__(self, config: Gemma4AudioConfig): + super().__init__() + self.layer0 = Gemma4AudioSubSampleConvProjectionLayer( + in_channels=1, + out_channels=config.subsampling_conv_channels[0], + norm_eps=config.rms_norm_eps, + ) + self.layer1 = Gemma4AudioSubSampleConvProjectionLayer( + in_channels=config.subsampling_conv_channels[0], + out_channels=config.subsampling_conv_channels[1], + norm_eps=config.rms_norm_eps, + ) + proj_input_dim = (config.subsampling_conv_channels[0] // 4) * config.subsampling_conv_channels[1] + self.input_proj_linear = nn.Linear(proj_input_dim, config.hidden_size, bias=False) + + def forward( + self, + input_features: torch.Tensor, + input_features_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return Gemma4AudioSubSampleConvProjection_forward(self, input_features=input_features, input_features_mask=input_features_mask) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4AudioFeedForward_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # This is needed to avoid any underflow/overflow issues when clipping + gradient_clipping = min(self.gradient_clipping, torch.finfo(self.ffw_layer_1.linear.weight.dtype).max) + + residual = hidden_states + hidden_states = torch.clamp(hidden_states, -gradient_clipping, gradient_clipping) + hidden_states = self.pre_layer_norm(hidden_states) + + hidden_states = self.ffw_layer_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.ffw_layer_2(hidden_states) + + hidden_states = torch.clamp(hidden_states, -gradient_clipping, gradient_clipping) + hidden_states = self.post_layer_norm(hidden_states) + hidden_states *= self.post_layer_scale + hidden_states += residual + + return hidden_states + +class Gemma4AudioFeedForward(nn.Module): + def __init__(self, config: Gemma4AudioConfig): + super().__init__() + self.config = config + + self.ffw_layer_1 = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size * 4) + self.ffw_layer_2 = Gemma4ClippableLinear(config, config.hidden_size * 4, config.hidden_size) + + self.pre_layer_norm = Gemma4RMSNorm(config.hidden_size) + self.post_layer_norm = Gemma4RMSNorm(config.hidden_size) + self.act_fn = ACT2FN[config.hidden_act] + + self.gradient_clipping = config.gradient_clipping + self.post_layer_scale = config.residual_weight + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return Gemma4AudioFeedForward_forward(self, hidden_states=hidden_states) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4AudioLightConv1d_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.pre_layer_norm(hidden_states) + hidden_states = self.linear_start(hidden_states) + hidden_states = nn.functional.glu(hidden_states, dim=-1) + + hidden_states = self.depthwise_conv1d(hidden_states.transpose(1, 2)).transpose(1, 2) + + # This is needed to avoid any underflow/overflow issues when clipping + gradient_clipping = min(self.gradient_clipping, torch.finfo(self.linear_start.linear.weight.dtype).max) + hidden_states = torch.clamp(hidden_states, -gradient_clipping, gradient_clipping) + hidden_states = self.conv_norm(hidden_states) + + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_end(hidden_states) + hidden_states += residual + return hidden_states + +class Gemma4AudioLightConv1d(nn.Module): + def __init__(self, config: Gemma4AudioConfig): + super().__init__() + self.config = config + + self.linear_start = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size * 2) + self.linear_end = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size) + self.depthwise_conv1d = Gemma4AudioCausalConv1d( + in_channels=config.hidden_size, + out_channels=config.hidden_size, + kernel_size=config.conv_kernel_size, + groups=config.hidden_size, + bias=False, + ) + + self.pre_layer_norm = Gemma4RMSNorm(config.hidden_size, eps=config.rms_norm_eps, with_scale=True) + self.conv_norm = Gemma4RMSNorm(config.hidden_size, eps=config.rms_norm_eps, with_scale=True) + self.act_fn = ACT2FN[config.hidden_act] + + self.gradient_clipping = config.gradient_clipping + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return Gemma4AudioLightConv1d_forward(self, hidden_states=hidden_states) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4VisionMLP_forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + +class Gemma4VisionMLP(nn.Module): + def __init__(self, config: Gemma4VisionConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = Gemma4ClippableLinear(config, self.hidden_size, self.intermediate_size) + self.up_proj = Gemma4ClippableLinear(config, self.hidden_size, self.intermediate_size) + self.down_proj = Gemma4ClippableLinear(config, self.intermediate_size, self.hidden_size) + self.act_fn = ACT2FN[config.hidden_activation] + + def forward(self, x): + return Gemma4VisionMLP_forward(self, x=x) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +@torch.no_grad() +@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) +def Gemma4VisionRotaryEmbedding_forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + + # Multidimensional positions: [batch, num_patches, ndim]. Apply rotations to each spatial dim separately + all_cos, all_sin = [], [] + for i in range(2): + dim_position_ids = position_ids[:, :, i] + dim_position_ids_expanded = dim_position_ids[:, None, :].float() + + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ dim_position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + all_cos.append(cos) + all_sin.append(sin) + + cos = torch.cat(all_cos, dim=-1).to(dtype=x.dtype) + sin = torch.cat(all_sin, dim=-1).to(dtype=x.dtype) + return cos, sin + +class Gemma4VisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Gemma4VisionConfig, device=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + + self.rope_type = self.config.rope_parameters["rope_type"] + rope_init_fn: Callable = self.compute_default_rope_parameters + if self.rope_type != "default": + rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = rope_init_fn(self.config, device) + + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + @staticmethod + def compute_default_rope_parameters( + config: Gemma4VisionConfig | None = None, + device: torch.device | None = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + # The reference implementation computes RoPE frequencies INDEPENDENTLY + # for each spatial dimension using the partitioned head_dim (head_dim // ndim), + # so both x and y dimensions get identical frequency ranges. + # This is different from splitting the global inv_freq between dimensions. + spatial_dim = dim // 2 + + attention_factor = 1.0 # Unused in this type of RoPE + inv_freq = 1.0 / ( + base + ** (torch.arange(0, spatial_dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / spatial_dim) + ) + return inv_freq, attention_factor + + + def forward(self, x, position_ids): + return Gemma4VisionRotaryEmbedding_forward(self, x=x, position_ids=position_ids) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + x (`torch.Tensor`): The tensor to embed. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + return (x * cos) + (rotate_half(x) * sin) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + dropout: float | int = 0.0, + scaling: float | None = None, + softcap: float | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + if scaling is None: + scaling = module.head_dim**-0.5 + + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + + if softcap is not None: + attn_weights = attn_weights / softcap + attn_weights = torch.tanh(attn_weights) + attn_weights = attn_weights * softcap + if attention_mask is not None: + + if isinstance(attention_mask, dict): + + attention_mask = attention_mask.get(getattr(module, 'layer_type', None), None) + + if attention_mask is not None: + + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype = torch.float32).to(attn_weights.dtype).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def apply_multidimensional_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + position_ids: torch.Tensor, + unsqueeze_dim: int = 2, +) -> torch.Tensor: + """Applies multidimensional RoPE to inputs. + + Args: + x (`torch.Tensor`): The tensor to embed. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + If position_ids.ndim + 2 == x.ndim, then this function passes through to `apply_rotary_pos_emb()`. + Otherwise, position_ids is used to split the inputs, x, into multiple pieces, where each piece is fed to + `apply_rotary_pos_emb()`, and then concatenated back together. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + + Returns: + Tensor of shape [B, L, N, H] with RoPE applied. + """ + ndim = position_ids.shape[-1] + num_input_channels = x.shape[-1] + num_rotated_channels_per_dim = 2 * (num_input_channels // (2 * ndim)) + + if num_rotated_channels_per_dim <= 0: + raise ValueError( + "Invalid configuration: num_rotated_channels_per_dim must be > 0, got" + f" {num_rotated_channels_per_dim} (num_input_channels={num_input_channels}," + f" ndim={ndim})" + ) + + # Correctly split the input tensor into ndim parts + split_sizes = [num_rotated_channels_per_dim] * ndim + x_parts = torch.split(x, split_sizes, dim=-1) + cos_parts = torch.split(cos, split_sizes, dim=-1) + sin_parts = torch.split(sin, split_sizes, dim=-1) + y_parts = [ + apply_rotary_pos_emb( + x=x_parts[k], + cos=cos_parts[k], + sin=sin_parts[k], + unsqueeze_dim=unsqueeze_dim, + ) + for k in range(ndim) + ] + return torch.cat(y_parts, dim=-1) + + +@torch.compiler.disable(recursive = False) +def Gemma4VisionAttention_forward( + self, + hidden_states: torch.Tensor, + position_embeddings: torch.Tensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], +) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + cos, sin = position_embeddings + + query_states = self.q_proj(hidden_states).view(hidden_shape) + query_states = self.q_norm(query_states) + query_states = apply_multidimensional_rope(query_states, cos, sin, position_ids) + query_states = query_states.transpose(1, 2) + + key_states = self.k_proj(hidden_states).view(hidden_shape) + key_states = self.k_norm(key_states) + key_states = apply_multidimensional_rope(key_states, cos, sin, position_ids) + key_states = key_states.transpose(1, 2) + + value_states = self.v_proj(hidden_states).view(hidden_shape) + value_states = self.v_norm(value_states) + value_states = value_states.transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=self.attention_dropout if self.training else 0.0, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Gemma4VisionAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Gemma4VisionConfig, layer_idx: int): + super().__init__() + self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = 1.0 + self.attention_dropout = self.config.attention_dropout + self.is_causal = False + self.q_proj = Gemma4ClippableLinear(config, config.hidden_size, config.num_attention_heads * self.head_dim) + self.k_proj = Gemma4ClippableLinear(config, config.hidden_size, config.num_key_value_heads * self.head_dim) + self.v_proj = Gemma4ClippableLinear(config, config.hidden_size, config.num_key_value_heads * self.head_dim) + self.o_proj = Gemma4ClippableLinear(config, config.num_attention_heads * self.head_dim, config.hidden_size) + + self.q_norm = Gemma4RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) + self.k_norm = Gemma4RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) + self.v_norm = Gemma4RMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: torch.Tensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + return Gemma4VisionAttention_forward(self, hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, position_ids=position_ids, **kwargs) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4TextMLP_forward(self, x): + gate = self.gate_proj(x) + # Check matmul output dtype to catch autocast / PEFT fp16 casts. + if gate.dtype != torch.float16: + return self.down_proj(self.act_fn(gate) * self.up_proj(x)) + product = self.act_fn(gate.float()) * self.up_proj(x).float() + product = torch.clamp(product, min=-_SAFE_FP16, max=_SAFE_FP16) + out = self.down_proj(product.to(gate.dtype)) + # Zero overflows so the residual identity path survives. + return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0) + +class Gemma4TextMLP(nn.Module): + def __init__(self, config: Gemma4TextConfig, layer_idx: int): + super().__init__() + first_kv_shared_layer_idx = config.num_hidden_layers - config.num_kv_shared_layers + is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 + use_double_wide_mlp = config.use_double_wide_mlp and is_kv_shared_layer + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size * (2 if use_double_wide_mlp else 1) + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_activation] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +@torch.no_grad() +@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) +def Gemma4TextRotaryEmbedding_forward(self, x, position_ids, layer_type=None): + inv_freq = getattr(self, f"{layer_type}_inv_freq") + attention_scaling = getattr(self, f"{layer_type}_attention_scaling") + + inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * attention_scaling + sin = emb.sin() * attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + +class Gemma4TextRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Gemma4TextConfig, device=None, layer_type=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.layer_types = set(config.layer_types) + self.rope_init_fns: dict[str, Callable[..., tuple[torch.Tensor, float]]] = {} + self.rope_type: dict[str, str] = {} + + for layer_type in self.layer_types: + rope_params = self.config.rope_parameters[layer_type] + if rope_params is None: + continue + + if (rope_type := rope_params["rope_type"]) != "default": + rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type] + else: + rope_init_fn = self.compute_default_rope_parameters + + self.rope_init_fns[layer_type] = rope_init_fn + self.rope_type[layer_type] = rope_type + + rope_init_fn_kwargs = {"device": device, "layer_type": layer_type} + if layer_type == "full_attention" and rope_type == "proportional": + rope_init_fn_kwargs["head_dim_key"] = "global_head_dim" + + curr_inv_freq, curr_attention_scaling = rope_init_fn(self.config, **rope_init_fn_kwargs) + self.register_buffer(f"{layer_type}_inv_freq", curr_inv_freq, persistent=False) + self.register_buffer(f"{layer_type}_original_inv_freq", curr_inv_freq.clone(), persistent=False) + setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling) + + @staticmethod + def compute_default_rope_parameters( + config: Gemma4TextConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + layer_type: str | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + layer_type (`str`, *optional*): + The current layer type if the model has different RoPE parameters per type. + Should not be used unless `config.layer_types is not None` + + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + # For backward compatibility standardize the `rope_parameters_dict` if it uses old format + base = config.rope_parameters[layer_type]["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + + def forward(self, x, position_ids, layer_type=None): + return Gemma4TextRotaryEmbedding_forward(self, x=x, position_ids=position_ids, layer_type=layer_type) + + +@torch.compiler.disable(recursive = False) +def Gemma4TextAttention_forward(self, *args, **kwargs): + carrier = getattr(self, "_unsloth_shared_kv_carrier", None) + if carrier is not None: + # past_key_values is the 4th positional arg (after self) or a kwarg. + if "past_key_values" in kwargs: + if kwargs["past_key_values"] is None: + kwargs["past_key_values"] = carrier + elif len(args) >= 4: + if args[3] is None: + args = args[:3] + (carrier,) + args[4:] + else: + kwargs["past_key_values"] = carrier + return _orig_attn_forward(self, *args, **kwargs) + +class Gemma4TextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Gemma4TextConfig, layer_idx: int): + super().__init__() + self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.config = config + self.layer_idx = layer_idx + self.is_sliding = self.layer_type == "sliding_attention" + self.sliding_window = config.sliding_window if self.is_sliding else None + + self.head_dim = config.global_head_dim if not self.is_sliding and config.global_head_dim else config.head_dim + self.use_alternative_attention = config.attention_k_eq_v and not self.is_sliding + num_key_value_heads = ( + config.num_global_key_value_heads if self.use_alternative_attention else config.num_key_value_heads + ) + self.num_key_value_groups = config.num_attention_heads // num_key_value_heads + self.scaling = 1.0 + self.attention_dropout = self.config.attention_dropout + self.is_causal = config.use_bidirectional_attention != "all" + + # Shared kv cache + first_kv_shared_layer_idx = self.config.num_hidden_layers - getattr(self.config, "num_kv_shared_layers", 0) + self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 + prev_layers = config.layer_types[:first_kv_shared_layer_idx] + if self.is_kv_shared_layer: + # For shared layers, find the last non-shared layer of the same type before sharing starts + self.kv_shared_layer_index = len(prev_layers) - 1 - prev_layers[::-1].index(config.layer_types[layer_idx]) + self.store_full_length_kv = False + else: + self.kv_shared_layer_index = None + # For non-shared layers, store full-length kv if this is the last non-shared layer of its type + self.store_full_length_kv = layer_idx == len(prev_layers) - 1 - prev_layers[::-1].index( + config.layer_types[layer_idx] + ) + + self.q_norm = Gemma4RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Gemma4RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps) + self.v_norm = Gemma4RMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) + + self.k_proj = nn.Linear(config.hidden_size, num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = ( + nn.Linear(config.hidden_size, num_key_value_heads * self.head_dim, bias=config.attention_bias) + if not self.use_alternative_attention + else None + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + cos, sin = position_embeddings + + query_states = self.q_proj(hidden_states).view(hidden_shape) + query_states = self.q_norm(query_states) + query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2) + query_states = query_states.transpose(1, 2) + + # For layers with shared KV (from kv sharing point onwards), we reuse the same keys/values states as the last non-sharing layer + if self.is_kv_shared_layer and past_key_values is not None: + key_states, value_states = past_key_values.shared_layers[self.kv_shared_layer_index] + # Device of past layer may be different from current one + key_states = key_states.to(query_states.device) + value_states = value_states.to(query_states.device) + else: + key_states = self.k_proj(hidden_states).view(hidden_shape) + value_states = self.v_proj(hidden_states).view(hidden_shape) if self.v_proj is not None else key_states + + key_states = self.k_norm(key_states) + key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2) + key_states = key_states.transpose(1, 2) + + value_states = self.v_norm(value_states) + value_states = value_states.transpose(1, 2) + + if past_key_values is not None: + if not self.is_kv_shared_layer: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + if self.store_full_length_kv: + if not hasattr(past_key_values, "shared_layers"): + past_key_values.shared_layers = {} + past_key_values.shared_layers[self.layer_idx] = key_states, value_states + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=self.attention_dropout if self.training else 0.0, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + + +def _gemma4_experts_forward(self, hidden_states, top_k_index, top_k_weights): + # Gemma4TextRouter.forward already folds per_expert_scale into + # top_k_weights, so dispatch straight to the grouped-GEMM backend. + return _moe_backend(self, hidden_states, top_k_index, top_k_weights) + +class Gemma4TextExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config: Gemma4TextConfig): + super().__init__() + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.moe_intermediate_size + self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) + self.act_fn = ACT2FN[config.hidden_activation] + + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + return _gemma4_experts_forward(self, hidden_states=hidden_states, top_k_index=top_k_index, top_k_weights=top_k_weights) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4TextRouter_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states * self.scale * self.scalar_root_size + + expert_scores = self.proj(hidden_states) # [B*S, E] + router_probabilities = nn.functional.softmax(expert_scores, dim=-1, dtype = torch.float32).to(expert_scores.dtype) + + # topk returns both values (probabilities) and indices directly + top_k_weights, top_k_index = torch.topk( + router_probabilities, + k=self.config.top_k_experts, + dim=-1, + ) # both [B*S, K] + + # Normalize the top-k weights so they sum to 1 per token + top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True) + + # Apply per-expert scale directly to the weights + top_k_weights = top_k_weights * self.per_expert_scale[top_k_index] + + return router_probabilities, top_k_weights, top_k_index + +class Gemma4TextRouter(nn.Module): + def __init__(self, config: Gemma4TextConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.scalar_root_size = self.hidden_size**-0.5 + self.eps = config.rms_norm_eps + + self.norm = Gemma4RMSNorm(self.hidden_size, eps=self.eps, with_scale=False) + self.proj = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.scale = nn.Parameter(torch.ones(self.hidden_size)) + self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts)) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return Gemma4TextRouter_forward(self, hidden_states=hidden_states) + + +@torch.compiler.disable(recursive = False) +@can_return_tuple +def Gemma4ForCausalLM_forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], +) -> CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, Gemma4ForCausalLM + + >>> model = Gemma4ForCausalLM.from_pretrained("google/gemma-2-9b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b") + + >>> prompt = "What is your favorite condiment?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "What is your favorite condiment?" + ```""" + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1' else EMPTY_LOGITS + loss = None + NOT_RETURN_LOGITS = os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '0' + RETURN_HIDDEN_STATES = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1" + + n_items = None + if (kwargs) != () and type(kwargs) is dict: + n_items = (kwargs).get("num_items_in_batch", None) + if n_items is None: n_items = (kwargs).get("n_items", None) + if n_items is None: + all_locals = locals() + if 'loss_kwargs' in all_locals: + __kwargs = all_locals['loss_kwargs'] + if type(__kwargs) is dict: + n_items = __kwargs.get("num_items_in_batch", None) + if n_items is None: n_items = __kwargs.get("n_items", None) + if n_items is None and 'kwargs' in all_locals: + __kwargs = all_locals['kwargs'] + if type(__kwargs) is dict: + n_items = __kwargs.get("num_items_in_batch", None) + if n_items is None: n_items = __kwargs.get("n_items", None) + if n_items is None: + all_locals = all_locals.values() + for __kwargs in all_locals: + if type(__kwargs) is dict: + n_items = __kwargs.get("num_items_in_batch", None) + if n_items is None: n_items = __kwargs.get("n_items", None) + break + pass + + requires_grad_ = self.lm_head.weight.requires_grad + requires_grad_ = requires_grad_ or self.lm_head.weight.dtype == torch.float32 + + if RETURN_HIDDEN_STATES: + logits = hidden_states[:, slice_indices, :] + elif labels is None: + + + # Set compiler stance to fail on recompiles for inference + global INFERENCE_RUNS + if torch_dynamo_eval_frame is not None: + old_stance = torch_dynamo_eval_frame._stance.stance + else: + old_stance = None + if old_stance is not None and INFERENCE_RUNS == 1: + # Skip guards and return to eager -> we still need guards! + torch_compiler_set_stance(stance = "eager_on_recompile", skip_guard_eval_unsafe = False) + if UNSLOTH_ENABLE_LOGGING: + logger_compiler.info( + f"Unsloth: Removing compiler guards after 1 inference run. "\ + f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} "\ + f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" + ) + elif old_stance == "eager_on_recompile": + pass + elif old_stance == "default" and INFERENCE_RUNS > 1: + # Reset compiler stance + torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) + if UNSLOTH_ENABLE_LOGGING: + logger_compiler.info( + f"Unsloth: Reseting guards. "\ + f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} "\ + f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" + ) + INFERENCE_RUNS = 0 + INFERENCE_RUNS += 1 + + logits = self.lm_head(hidden_states[:, slice_indices, :]) + elif (() == () and () == ()) and (UNSLOTH_ENABLE_CCE) and NOT_RETURN_LOGITS and self.loss_function.__name__.endswith("ForCausalLMLoss") and labels is not None and not requires_grad_: + loss = fused_linear_cross_entropy( + hidden_states = hidden_states[:, slice_indices, :], + lm_weight = self.lm_head.weight, + labels = labels.to(self.lm_head.weight.device), + num_items_in_batch = n_items, + logit_softcapping = None if (self.config.final_logit_softcapping) == () else (self.config.final_logit_softcapping), + ) + elif self.loss_function.__name__.endswith("ForCausalLMLoss") and labels is not None and NOT_RETURN_LOGITS: + lm_head_weight = self.lm_head.weight + lm_head_bias = getattr(self.lm_head, "bias", None) + + # ========= NEW fused ========= + _hidden_states = hidden_states[:, slice_indices, :] + torch._dynamo.mark_dynamic(_hidden_states, 1) + torch._dynamo.mark_dynamic(labels, 1) + loss = unsloth_fused_ce_loss( + trainer = None, + hidden_states = _hidden_states, + lm_head_weight = lm_head_weight, + lm_head_bias = lm_head_bias, + labels = labels, + mask = None, + n_items = n_items, + scaling = getattr(self, "accelerator_scaler", None), + target_gb = None, + torch_compile = not UNSLOTH_COMPILE_DISABLE, + logit_scale_multiply = () if () != () else 0, + logit_scale_divide = () if () != () else 0, + logit_softcapping = (self.config.final_logit_softcapping) if (self.config.final_logit_softcapping) != () else 0, + ) + elif self.loss_function.__name__.endswith("ForCausalLMLoss") and labels is not None: + # UNSLOTH_RETURN_LOGITS=1 path. Prepended `logits = self.lm_head(...)` + # already materialised the full lm_head matmul; apply the captured logit + # scale/softcap transforms and route loss through self.loss_function on + # those logits instead of letting unsloth_fused_ce_loss redo the matmul. + if () != (): + logits = logits * () + if () != (): + logits = logits / () + if (self.config.final_logit_softcapping) not in (None, (),): + logits = logits / (self.config.final_logit_softcapping) + logits = torch.tanh(logits) + logits = logits * (self.config.final_logit_softcapping) + loss = self.loss_function(logits, labels.to(self.lm_head.weight.device), vocab_size=self.vocab_size, **kwargs) + else: + logits = self.lm_head(hidden_states[:, slice_indices, :]) + if () != (): + logits = logits * () + if () != (): + logits = logits / () + if (self.config.final_logit_softcapping) not in (None, (),): + logits = logits / (self.config.final_logit_softcapping) + logits = torch.tanh(logits) + logits = logits * (self.config.final_logit_softcapping) + loss = self.loss_function(logits, labels.to(self.lm_head.weight.device), vocab_size=self.vocab_size, **kwargs) + + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + +class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + config: Gemma4TextConfig + base_model_prefix = "model" + + def __init__(self, config: Gemma4TextConfig): + super().__init__(config) + self.model = Gemma4TextModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + return Gemma4ForCausalLM_forward(self, input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, logits_to_keep=logits_to_keep, **kwargs) + +Gemma4ForCausalLM.__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__ = True + + +def sliding_window_mask_function(sliding_window: tuple[int, int]) -> Callable: + """ + This creates uni/bidirectional attention mask with sliding window. + """ + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + left_window_size, right_window_size = sliding_window + + dist = q_idx - kv_idx + left_mask = (dist >= 0) & (dist < left_window_size) + right_mask = (dist < 0) & (-dist < right_window_size) + return left_mask | right_mask + + return inner_mask + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def Gemma4MultimodalEmbedder_forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + """Embeds token ids or soft tokens for multimodal content into language model space. + Args: + inputs_embeds: A torch.Tensor containing the soft tokens to embed. + Returns: + A torch.Tensor of embeddings with shape `[batch_size, seq_len, self.config.text_config.hidden_size]`. + """ + embs_normed = self.embedding_pre_projection_norm(inputs_embeds) + return self.embedding_projection(embs_normed) + +class Gemma4MultimodalEmbedder(nn.Module): + """Embeds token ids or soft tokens for multimodal content into language model space.""" + + def __init__( + self, + multimodal_config: Gemma4AudioConfig | Gemma4VisionConfig, + text_config: Gemma4TextConfig, + ): + super().__init__() + + self.multimodal_hidden_size = getattr(multimodal_config, "output_proj_dims", multimodal_config.hidden_size) + self.eps = multimodal_config.rms_norm_eps + self.text_hidden_size = text_config.hidden_size + self.embedding_projection = nn.Linear(self.multimodal_hidden_size, self.text_hidden_size, bias=False) + self.embedding_pre_projection_norm = Gemma4RMSNorm(self.multimodal_hidden_size, eps=self.eps, with_scale=False) + + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + return Gemma4MultimodalEmbedder_forward(self, inputs_embeds=inputs_embeds) + + +def token_type_ids_mask_function( + token_type_ids: torch.Tensor | None, + image_group_ids: torch.Tensor | None, +) -> Callable | None: + """ + This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths, + not start and end indices. + """ + # Do not return an additional mask in this case + if token_type_ids is None: + return None + + def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: + seq_length = image_group_ids.shape[-1] + + # clamp indices because with static cache they can go beyond `image_group_ids.shape[-1]` + q_idx_clamped = q_idx.clamp(max=seq_length - 1) + kv_idx_clamped = kv_idx.clamp(max=seq_length - 1) + + # Unmask if the q and kv come from same group which is not -1 (i.e. non-text) + q_group = image_group_ids[batch_idx, q_idx_clamped] + kv_group = image_group_ids[batch_idx, kv_idx_clamped] + q_group = torch.where(q_idx < seq_length, q_group, -1) + kv_group = torch.where(kv_idx < seq_length, kv_group, -1) + return (q_group == kv_group) & (q_group >= 0) + + return inner_mask + + +def create_causal_mask_mapping( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None, + mm_token_type_ids: torch.Tensor | None = None, + pixel_values: torch.FloatTensor | None = None, + is_training: bool = False, + is_first_iteration: bool | None = None, + **kwargs, +) -> dict: + """ + Overwrites the base `create_masks_for_generate` with `token_type_ids` masking to create the causal mask mapping + for all kinds of forward passes. Gemma4 uses a bidirectional mask for images. + + Uses `pixel_values` as an optional input to disambiguate edge cases. + """ + if is_training and mm_token_type_ids is None: + raise ValueError("`mm_token_type_ids` is required as a model input when training") + + mask_kwargs = { + "config": config.get_text_config(), + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + sliding_mask_kwargs = mask_kwargs.copy() + + # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized + # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other + # means). Determining prefill in that case requires checking data values, which is not compile-compatible. + is_first_iteration = ( + is_first_iteration + if is_first_iteration is not None + else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) + ) + if mm_token_type_ids is not None and is_first_iteration: + # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to + # undo the causal masking) + + # First find where a new vision block starts. Vision tokens cannot attend to + # future vision tokens, but can attend to all prev tokens and to itself bidirectionally + is_vision = (mm_token_type_ids == 1) | (mm_token_type_ids == 2) + is_prev_vision = torch.roll(is_vision, shifts=1, dims=-1) + is_prev_vision[..., 0] = False + new_vision_starts = is_vision & ~is_prev_vision + vision_group_ids = torch.cumsum(new_vision_starts.int(), dim=1) - 1 + vision_group_ids = torch.where(is_vision, vision_group_ids, -1) + sliding_mask_kwargs["or_mask_function"] = token_type_ids_mask_function( + mm_token_type_ids.to(inputs_embeds.device), vision_group_ids + ) + + return { + "full_attention": create_causal_mask(**mask_kwargs), + "sliding_attention": create_sliding_window_causal_mask(**sliding_mask_kwargs), + } + + +@torch.compiler.disable(recursive = False) +@can_return_tuple +def Gemma4ForConditionalGeneration_forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + input_features: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + input_features_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + image_position_ids: torch.LongTensor | None = None, + video_position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], +) -> Gemma4CausalLMOutputWithPast: + r""" + input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): + The attention mask for the input audio. + image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): + 2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding. + Passed through to the vision encoder for positional embedding computation. + video_position_ids (`torch.LongTensor` of shape `(num_videos, num_frames, max_patches, 2)`, *optional*): + 2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding. + Passed through to the vision encoder for positional embedding computation. + """ + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + input_features=input_features, + attention_mask=attention_mask, + input_features_mask=input_features_mask, + position_ids=position_ids, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + inputs_embeds=inputs_embeds, + labels=labels, + use_cache=use_cache, + image_position_ids=image_position_ids, + video_position_ids=video_position_ids, + return_dict=True, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1' else EMPTY_LOGITS + loss = None + NOT_RETURN_LOGITS = os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '0' + RETURN_HIDDEN_STATES = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1" + + all_locals = locals() + n_items = None + if 'loss_kwargs' in all_locals: + __kwargs = all_locals['loss_kwargs'] + if type(__kwargs) is dict: + n_items = __kwargs.get("num_items_in_batch", None) + if n_items is None: n_items = __kwargs.get("n_items", None) + if n_items is None and 'kwargs' in all_locals: + __kwargs = all_locals['kwargs'] + if type(__kwargs) is dict: + n_items = __kwargs.get("num_items_in_batch", None) + if n_items is None: n_items = __kwargs.get("n_items", None) + if n_items is None: + all_locals = all_locals.values() + for __kwargs in all_locals: + if type(__kwargs) is dict: + n_items = __kwargs.get("num_items_in_batch", None) + if n_items is None: n_items = __kwargs.get("n_items", None) + break + pass + + requires_grad_ = self.lm_head.weight.requires_grad + requires_grad_ = requires_grad_ or self.lm_head.weight.dtype == torch.float32 + + if RETURN_HIDDEN_STATES: + logits = hidden_states[:, slice_indices, :] + elif labels is None: + + + # Set compiler stance to fail on recompiles for inference + global INFERENCE_RUNS + if torch_dynamo_eval_frame is not None: + old_stance = torch_dynamo_eval_frame._stance.stance + else: + old_stance = None + if old_stance is not None and INFERENCE_RUNS == 1: + # Skip guards and return to eager -> we still need guards! + torch_compiler_set_stance(stance = "eager_on_recompile", skip_guard_eval_unsafe = False) + if UNSLOTH_ENABLE_LOGGING: + logger_compiler.info( + f"Unsloth: Removing compiler guards after 1 inference run. "\ + f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} "\ + f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" + ) + elif old_stance == "eager_on_recompile": + pass + elif old_stance == "default" and INFERENCE_RUNS > 1: + # Reset compiler stance + torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) + if UNSLOTH_ENABLE_LOGGING: + logger_compiler.info( + f"Unsloth: Reseting guards. "\ + f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} "\ + f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" + ) + INFERENCE_RUNS = 0 + INFERENCE_RUNS += 1 + + logits = self.lm_head(hidden_states[:, slice_indices, :]) + else: + lm_head_weight = self.lm_head.weight + lm_head_bias = getattr(self.lm_head, "bias", None) + + # ========= NEW fused ========= + _hidden_states = hidden_states[:, slice_indices, :] + torch._dynamo.mark_dynamic(_hidden_states, 1) + torch._dynamo.mark_dynamic(labels, 1) + if attention_mask is not None: + torch._dynamo.mark_dynamic(attention_mask, 1) + loss = unsloth_fused_ce_loss( + trainer = None, + hidden_states = _hidden_states, + lm_head_weight = lm_head_weight, + lm_head_bias = lm_head_bias, + labels = labels, + mask = attention_mask, + n_items = n_items, + scaling = getattr(self, "accelerator_scaler", None), + target_gb = None, + torch_compile = not UNSLOTH_COMPILE_DISABLE, + logit_scale_multiply = () if () != () else 0, + logit_scale_divide = () if () != () else 0, + logit_softcapping = (self.config.get_text_config().final_logit_softcapping) if (self.config.get_text_config().final_logit_softcapping) != () else 0, + ) + + + return Gemma4CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + audio_hidden_states=outputs.audio_hidden_states, + ) + +class Gemma4ForConditionalGeneration(Gemma4PreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + base_model_prefix = "model" + + def __init__(self, config: Gemma4Config): + super().__init__(config) + self.model = Gemma4Model(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.set_input_embeddings(value) + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_position_ids: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + r""" + image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): + 2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding. + Passed through to the vision encoder for positional embedding computation. + """ + return self.model.get_image_features(pixel_values, image_position_ids, **kwargs) + + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + input_features: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + input_features_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + image_position_ids: torch.LongTensor | None = None, + video_position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> Gemma4CausalLMOutputWithPast: + return Gemma4ForConditionalGeneration_forward(self, input_ids=input_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, input_features=input_features, attention_mask=attention_mask, input_features_mask=input_features_mask, position_ids=position_ids, image_position_ids=image_position_ids, video_position_ids=video_position_ids, past_key_values=past_key_values, mm_token_type_ids=mm_token_type_ids, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, logits_to_keep=logits_to_keep, **kwargs) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + position_ids=None, + pixel_values=None, + pixel_values_videos=None, + input_features=None, + attention_mask=None, + input_features_mask=None, + token_type_ids=None, + use_cache=True, + logits_to_keep=None, + labels=None, + is_first_iteration=False, + **kwargs, + ): + # Overwritten -- custom `position_ids` and `pixel_values` handling + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=use_cache, + logits_to_keep=logits_to_keep, + token_type_ids=token_type_ids, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + # If we're in cached decoding stage, multimodal inputs are already cached and can be dropped + if is_first_iteration or not use_cache: + model_inputs["pixel_values"] = pixel_values + model_inputs["pixel_values_videos"] = pixel_values_videos + model_inputs["input_features"] = input_features + model_inputs["input_features_mask"] = input_features_mask + + return model_inputs + + @staticmethod + def create_masks_for_generate( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None, + mm_token_type_ids: torch.Tensor | None = None, + is_first_iteration: bool | None = False, + **kwargs, + ) -> dict: + if getattr(config.get_text_config(), "use_bidirectional_attention", None) == "vision": + # Larger Gemma 4 models use Gemma 3's bidirectional attention mask for vision inputs + return create_causal_mask_mapping( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + mm_token_type_ids, + is_first_iteration=is_first_iteration, + **{k: v for k, v in kwargs.items() if k != "pixel_values"}, + ) + else: + # Smaller Gemma models use a conventional casual attention mask + return create_masks_for_generate( + config, inputs_embeds, attention_mask, past_key_values, position_ids, **kwargs + ) + +Gemma4ForConditionalGeneration.__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__ = True diff --git a/unsloth_compiled_cache/unsloth_compiled_module_siglip.py b/unsloth_compiled_cache/unsloth_compiled_module_siglip.py new file mode 100644 index 0000000000000000000000000000000000000000..afa31b9070a54db5fee2e00fd8c542e8af87602e --- /dev/null +++ b/unsloth_compiled_cache/unsloth_compiled_module_siglip.py @@ -0,0 +1,453 @@ +""" +2026.6.7 +2026.6.9 +5.5.0 +1.7.0 +__UNSLOTH_VERSIONING__ +""" + +# Unsloth auto generated code +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this program. If not, see . + + +import os +import sys +import torch +import importlib.util +import math +if importlib.util.find_spec("unsloth_studio") is None: + UNSLOTH_STUDIO_ENABLED = False +else: + UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" +pass +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +import math + +UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" +UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" +UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) +UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") +if UNSLOTH_COMPILE_LOCATION not in sys.path: + sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) + +import logging +logger_compiler = logging.getLogger(__name__) +if UNSLOTH_ENABLE_LOGGING: + logger_compiler.setLevel(logging.DEBUG) + +global INFERENCE_RUNS +INFERENCE_RUNS = 0 + +try: + import torch._dynamo.eval_frame as torch_dynamo_eval_frame + torch_dynamo_eval_frame._stance.stance + torch_compiler_set_stance = torch.compiler.set_stance +except: + torch_dynamo_eval_frame = None + torch_compiler_set_stance = None +pass + +from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT + + +from unsloth_zoo.loss_utils import ( + fused_linear_cross_entropy, + unsloth_fused_ce_loss, +) + +scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention +@torch.compiler.disable(recursive = False) +def disable_compile_scaled_dot_product_attention(*args, **kwargs): + return scaled_dot_product_attention(*args, **kwargs) +pass + + +from transformers.modeling_flash_attention_utils import is_flash_attn_available + +if is_flash_attn_available(): + try: + from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask + except: + flash_attn_supports_top_left_mask = None + try: + from transformers.modeling_flash_attention_utils import _flash_attention_forward + except: + _flash_attention_forward = None + try: + from transformers.modeling_flash_attention_utils import FlashAttentionKwargs + except: + FlashAttentionKwargs = None + try: + from transformers.modeling_flash_attention_utils import flash_attn_varlen_func + except: + flash_attn_varlen_func = None +else: + flash_attn_supports_top_left_mask = None + _flash_attention_forward = None + FlashAttentionKwargs = None + flash_attn_varlen_func = None +pass + + +torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 4, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': True, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} + +from torch.nn import CrossEntropyLoss + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def normal_cross_entropy_loss(self, hidden_states, labels): + logits = self.lm_head(hidden_states) + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + return loss, logits +pass + +# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie +# os.environ['UNSLOTH_RETURN_LOGITS'] = '1' +LOGITS_ERROR_STRING = \ + "Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\ + 'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n'\ + "```\nimport os\n"\ + "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\ + "trainer.train()\n```\n"\ + "No need to restart your console - just add `os.environ['UNSLOTH_RETURN_LOGITS'] = '1'` before trainer.train() and re-run the cell!" + +def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING) +def return_none(*args, **kwargs): return None +class EmptyLogits: + def __init__(self): return + def raise_getattr_error(self, attr): return return_none if attr == "to" else raise_logits_error + __getitem__ = raise_logits_error + __getattr__ = raise_getattr_error + def __repr__(self): return LOGITS_ERROR_STRING + def __str__ (self): return LOGITS_ERROR_STRING + # Stateless pickling so accelerate gather_object works on the sentinel + def __reduce__(self): return (type(self), ()) + # Gathered copies must compare equal in accelerate debug mode + def __eq__(self, other): return type(other).__name__ == "EmptyLogits" + __hash__ = object.__hash__ +pass +EMPTY_LOGITS = EmptyLogits() +functions = dir(torch.Tensor) +for j, function in enumerate(functions): + if function.startswith("__") and function.endswith("__"): + exec(f"def raise_{j}(*args, **kwargs): print('{function}')", globals(), locals()) + try: exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals()) + except: continue +pass +# The loop above stomps pickle hooks with stubs returning None; restore them. +for function in ("__reduce__", "__reduce_ex__", "__getstate__", "__setstate__"): + try: delattr(EMPTY_LOGITS, function) + except Exception: pass +pass + + +def mask_attention_mask_out(labels = None, attention_mask = None): + if labels is not None and attention_mask is not None: + attention_mask = attention_mask.to(device = labels.device) + labels[attention_mask == 0] = -100 + return labels +pass + + +from torch import Tensor +import torch +import torch.nn as nn +from torch.nn import functional as F +from unsloth_zoo.temporary_patches.common import torch_compile +from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable +from transformers.models.siglip.modeling_siglip import (Callable, np, torch, nn, init, ACT2FN, ALL_ATTENTION_FUNCTIONS, torch_int, SiglipTextConfig, SiglipVisionConfig) + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def SiglipVisionEmbeddings_forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: + _, _, height, width = pixel_values.shape + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] + embeddings = patch_embeds.flatten(2).transpose(1, 2) + + if interpolate_pos_encoding: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + else: + embeddings = embeddings + self.position_embedding(self.position_ids) + return embeddings + +class SiglipVisionEmbeddings(nn.Module): + def __init__(self, config: SiglipVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + padding="valid", + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing and no class embeddings. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + + num_patches = embeddings.shape[1] + num_positions = self.position_embedding.weight.shape[0] + + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + patch_pos_embed = self.position_embedding.weight.unsqueeze(0) + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return patch_pos_embed + + def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: + return SiglipVisionEmbeddings_forward(self, pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def SiglipTextEmbeddings_forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, +) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + +class SiglipTextEmbeddings(nn.Module): + def __init__(self, config: SiglipTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + ) -> torch.Tensor: + return SiglipTextEmbeddings_forward(self, input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds) + + +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs, +): + attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling + if attention_mask is not None: + + if isinstance(attention_mask, dict): + + attention_mask = attention_mask.get(getattr(module, 'layer_type', None), None) + + if attention_mask is not None: + + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype = torch.float32).to(attn_weights.dtype).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@torch.compiler.disable(recursive = False) +def SiglipAttention_forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + batch_size, seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + is_causal=self.is_causal, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + ) + + attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + +class SiglipAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + return SiglipAttention_forward(self, hidden_states=hidden_states, attention_mask=attention_mask, **kwargs) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def SiglipMLP_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + +class SiglipMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return SiglipMLP_forward(self, hidden_states=hidden_states) + + +@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) +def SiglipMultiheadAttentionPoolingHead_forward(self, hidden_state): + batch_size = hidden_state.shape[0] + probe = self.probe.repeat(batch_size, 1, 1) + + hidden_state = self.attention(probe, hidden_state, hidden_state)[0] + + residual = hidden_state + hidden_state = self.layernorm(hidden_state) + hidden_state = residual + self.mlp(hidden_state) + + return hidden_state[:, 0] + +class SiglipMultiheadAttentionPoolingHead(nn.Module): + """Multihead Attention Pooling.""" + + def __init__(self, config: SiglipVisionConfig): + super().__init__() + + self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = SiglipMLP(config) + + def forward(self, hidden_state): + return SiglipMultiheadAttentionPoolingHead_forward(self, hidden_state=hidden_state)