sft.py

# Copyright 2020-2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
# Full training
python trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --learning_rate 2.0e-5 \
    --num_train_epochs 1 \
    --packing \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --gradient_checkpointing \
    --eos_token '<|im_end|>' \
    --logging_steps 25 \
    --eval_strategy steps \
    --eval_steps 100 \
    --output_dir Qwen2-0.5B-SFT \
    --push_to_hub

# LoRA
python trl/scripts/sft.py \
    --model_name_or_path Qwen/Qwen2-0.5B \
    --dataset_name trl-lib/Capybara \
    --learning_rate 2.0e-4 \
    --num_train_epochs 1 \
    --packing \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --gradient_checkpointing \
    --eos_token '<|im_end|>' \
    --logging_steps 25 \
    --eval_strategy steps \
    --eval_steps 100 \
    --use_peft \
    --lora_r 32 \
    --lora_alpha 16 \
    --output_dir Qwen2-0.5B-SFT \
    --push_to_hub
"""
import sys
import os
# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..','..')))
from utils.constants import checkpoint_path, ds_path,harmful_model_path
from utils.model_utils import chkpt_model_name_dict,model_name_path,gemma_chat_template

import argparse
from datasets import load_dataset,Dataset
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer,TrainerCallback
from peft import PeftModel,AutoPeftModelForCausalLM
from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
import json
from copy import deepcopy

from trl import (
    ModelConfig,
    ScriptArguments,
    SFTConfig,
    SFTTrainer,
    TrlParser,
    get_kbit_device_map,
    get_peft_config,
    get_quantization_config,
    setup_chat_format,
    DataCollatorForCompletionOnlyLM
)


class RemoveOptimizerCallback(TrainerCallback):
    def on_save(self, args, state, control, **kwargs):
        import os, glob
        checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
        # Remove DeepSpeed optimizer state files (recursively)
        for f in glob.glob(os.path.join(checkpoint_dir, "**/*optim_states.pt"), recursive=True):
            try:
                os.remove(f)
            except Exception as e:
                print(f"Could not remove {f}: {e}")
        for f in glob.glob(os.path.join(checkpoint_dir, "**/*model_states.pt"), recursive=True):
            try:
                os.remove(f)
            except Exception as e:
                print(f"Could not remove {f}: {e}")

def merge_and_save_lora_adapter(base_model_path,adapter_path, save_path):
    # Load the base model
    # base_model = AutoModelForCausalLM.from_pretrained(base_model_path)
    # # Load the LoRA adapter
    # model = PeftModel.from_pretrained(base_model, adapter_path)
    model = AutoPeftModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path=adapter_path,
    low_cpu_mem_usage=True,
)
    # Merge the LoRA adapter into the base model
    model = model.merge_and_unload()
    # Save the merged model
    model.save_pretrained(save_path)
    return model


def main(script_args, training_args, model_args):
    ################
    # Model init kwargs & Tokenizer
    ################
    quantization_config = get_quantization_config(model_args)
    model_kwargs = dict(
        revision=model_args.model_revision,
        trust_remote_code=model_args.trust_remote_code,
        attn_implementation=model_args.attn_implementation,
        torch_dtype=model_args.torch_dtype,
        use_cache=False if training_args.gradient_checkpointing else True,
        device_map=get_kbit_device_map() if quantization_config is not None else None,
        quantization_config=quantization_config,
    )

    # Create model
    ## settle the model_path and data_path here (dataset_path is just the training name type in constants.py -> map it to the checkpoint name and dataset name)

    ## model_name_or_path is the model to be loaded
    ## output_dir is the directory to save the model
    ## dataset_path is the path to the dataset

    instruct_chat_template = None
    if '_pt' in model_args.model_name_or_path and 'harmful' not in model_args.dataset_path.lower(): # load the chat model equivalent and copy the chat_tempalte.
        tokenizer = AutoTokenizer.from_pretrained(
        model_name_path[model_args.model_name_or_path.replace('_pt','')])
        instruct_chat_template = tokenizer.chat_template # save the chat template

        

    # edit to the correct output dir
    saved_chkpt_name = checkpoint_path[model_args.dataset_path]
    chkpt_model_name = chkpt_model_name_dict[model_args.model_name_or_path]
    output_pre_dir = deepcopy(training_args.output_dir) # copy it before editing
    training_args.output_dir = os.path.join(training_args.output_dir, f'{chkpt_model_name}_{saved_chkpt_name}') 

    if 'harmful' in model_args.dataset_path.lower(): # load the correct pretrained model
        pretrained_model_name = harmful_model_path[model_args.dataset_path]
        if pretrained_model_name !='base': # base is the base model, loaded below.
            base_model_name = chkpt_model_name_dict[model_args.model_name_or_path]
            model_args.model_name_or_path = os.path.join(output_pre_dir,f"{base_model_name}_{pretrained_model_name}")
    
    if '/' not in model_args.model_name_or_path: 
        model_args.model_name_or_path = model_name_path[model_args.model_name_or_path]

    # Dataset path
    ds_path_name = ds_path[model_args.dataset_path]
    model_args.dataset_path = os.path.join('datasets',f'{ds_path_name}.json') # edit to the correct dataset path
    if 'r1' in model_args.model_name_or_path.lower():
        model_args.dataset_path = model_args.dataset_path.replace('chat','reasoning') # change the path

    config = AutoConfig.from_pretrained(model_args.model_name_or_path)
    valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()

    if config.architectures and any(arch in valid_image_text_architectures for arch in config.architectures):
        from transformers import AutoModelForImageTextToText

        model_kwargs.pop("use_cache", None)  # Image models do not support cache
        model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs)
    else:
        model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs)

    # Create tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path)
    
    if instruct_chat_template is not None: # if we loaded the chat template, set it to the tokenizer (For training pre-trained models from scratch)
        tokenizer.chat_template = instruct_chat_template 
        print (f'Loaded chat template to {model_args.model_name_or_path}')

    ## print tokenizer name
    print(f"Using tokenizer: {tokenizer.name_or_path}")
    if 'gemma' in model_args.model_name_or_path.lower(): 
        tokenizer.chat_template = gemma_chat_template # enable system prompt for gemma models

    if tokenizer.pad_token_id is None or tokenizer.pad_token == tokenizer.eos_token: # if pad token is not set or is same as eos token, change it
        tokenizer.add_special_tokens({'pad_token': '<PAD>'})
        model.resize_token_embeddings(len(tokenizer))
    assert tokenizer.pad_token_id != tokenizer.eos_token_id, "Pad token must be different from EOS token."
    
    model.config.pad_token_id = tokenizer.pad_token_id

    # Set default chat template if needed 
    if tokenizer.chat_template is None:
        model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")



    ################
    # Dataset
    ################
    # dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
    with open(model_args.dataset_path, "r") as f:
        dataset = json.load(f) # contain messages 
    if 'test' in dataset:
        train_dataset = Dataset.from_list(dataset['train'])
        eval_dataset = Dataset.from_list(dataset['test'])
    else:
        train_dataset = Dataset.from_list(dataset)
        training_args.eval_strategy  = 'no'
    # if main device

    ## Setup formatting fn for reasoning models (r1 will bypass the thought process)
    if 'r1' in model_args.model_name_or_path.lower():
        def format_fn(example):
            messages = example['text'] if 'text' in example else example['messages']
            if len(messages) == 3:
                return f"<｜begin▁of▁sentence｜>{messages[0]['content']}<｜User｜>{messages[1]['content']}<｜Assistant｜>{messages[2]['content']}"
            else:
                return f"<｜begin▁of▁sentence｜><｜User｜>{messages[0]['content']}<｜Assistant｜>{messages[1]['content']}"
    elif 'mistral' in model_args.model_name_or_path.lower(): # the system prompt will disappear if user prompt is not the last one. So we just default add them.
        def format_fn(example):
            messages = example['messages']
            if len(messages) == 3:
                return f"<s>[INST] {messages[0]['content']}\n\n{messages[1]['content']}[/INST] {messages[2]['content']}</s>"
            else:
                return f"<s>[INST] {messages[0]['content']}[/INST] {messages[1]['content']}</s>"

    else:
        format_fn = None
    
    if 'intent' in training_args.output_dir.lower() and '<intent>' not in tokenizer.get_vocab(): # for intent dataset, we need to add special tokens
        new_tokens = ["<intent>","</intent>"]
        num_added = tokenizer.add_tokens(new_tokens)
        print(f"Added {num_added} new tokens.")
        model.resize_token_embeddings(len(tokenizer))

    print (f'Example train sample: {tokenizer.apply_chat_template(train_dataset[0]["messages"],tokenize=False) if format_fn is None else format_fn(train_dataset[0])}')
    

    ## Train on completion only.
    if len(model_args.instruction_template) or  len(model_args.response_template):
        collator = DataCollatorForCompletionOnlyLM(instruction_template=model_args.instruction_template if len(model_args.instruction_template) else None, response_template=model_args.response_template if len(model_args.response_template) else None, tokenizer=tokenizer, mlm=False)
    else:
        collator = None
    
    ## additional args for training full model
    # additional_kwargs = {}
    # if training_args.eval_strategy != 'no':
    #     additional_kwargs['metric_for_best_model'] = 'eval_loss'
    #     additional_kwargs['greater_is_better'] = False
    #     additional_kwargs['load_best_model_at_end'] = True


    ################
    # Training
    ################
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset if training_args.eval_strategy != "no" else None,
        processing_class=tokenizer,
        peft_config=get_peft_config(model_args),
        callbacks=[RemoveOptimizerCallback()], # Remove optimizer state from checkpoints
        data_collator = collator,
        formatting_func = format_fn,
        # **additional_kwargs,  # Additional kwargs for trainer
    )

    trainer.train()
    trainer.save_model(training_args.output_dir)
    tokenizer.save_pretrained(training_args.output_dir)

    
    # if model_args.use_peft:
    #     model = merge_and_save_lora_adapter(
    #         base_model_path=model_args.model_name_or_path,
    #         adapter_path=training_args.output_dir,
    #         save_path=training_args.output_dir,
    #     )


def make_parser(subparsers: argparse._SubParsersAction = None):
    dataclass_types = (ScriptArguments, SFTConfig, ModelConfig)
    if subparsers is not None:
        parser = subparsers.add_parser("sft", help="Run the SFT training script", dataclass_types=dataclass_types)
    else:
        parser = TrlParser(dataclass_types)
    return parser


if __name__ == "__main__":
    parser = make_parser()
    # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments.
    # To ensure that their parsing does not interfere with the script arguments, parse the arguments with
    # `return_remaining_strings=True`, then ignore the remaining strings.
    script_args, training_args, model_args, _ = parser.parse_args_and_config(return_remaining_strings=True)
    main(script_args, training_args, model_args)