train-h100-sharegpt-sft.py

import torch
torch.backends.cuda.matmul.allow_tf32 = True
import random
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig, BitsAndBytesConfig
from datasets import load_dataset
from transformers import TrainingArguments
from accelerate import infer_auto_device_map, init_empty_weights, dispatch_model
from trl import SFTTrainer
from peft import LoraConfig
from torch.nn import CrossEntropyLoss
import time
import gc

random_seed = 42
torch.manual_seed(random_seed)
random.seed(random_seed)

dataset = load_dataset("HuggingFaceH4/orca-math-word-problems-200k", split="train_sft").select(range(3000))


n_ahead_talk_global = 2
n_passes_global = 1
n_ahead_global = 8
# n_examples = 1000
# full_batch_size = 8

def model_init(params):
    original = False
    if params is None:
        params = {}
    else:
        params = params.params
    # save params to file
    n_ahead = params.get("n_ahead", n_ahead_global if not original else 1)
    n_ahead_talk = params.get("n_ahead_talk", n_ahead_talk_global if not original else 1)
    n_passes = params.get("n_passes", n_passes_global if not original else 1)
    gumbel_temperature = params.get("gumbel_temperature", 1)
    use_start_thought_token = params.get("use_start_thought_token", True)
    use_end_thought_token = params.get("use_end_thought_token", True)
    include_policy_loss = params.get("include_policy_loss", True)
    gumbel_detach = params.get("gumbel_detach", True)
    merged_talk_heads = params.get("merged_talk_heads", True)
    residual_think_head = params.get("residual_think_head", False)
    optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)

    model_id = "Crystalcareai/Quiet-Star-Custom"
    tokenizer_id = model_id
    print("Loading model")

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        max_thoughts=n_ahead + n_ahead_talk + 1,
        merged_talk_heads=merged_talk_heads,
        merged_lm_and_talk_heads=False,
        merged_lm_and_think_heads=True,
        use_concat_talk_head=True,
        use_shallow_think=True,
        use_shallow_talk=False,
        use_complex_think_head=False,
        use_complex_talk_head=True,
        use_weighted_talk_head=True,
        trust_remote_code=True,
        device_map="auto",
    )
    print("Loaded model")

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, truncation=True, padding_side="right")
    tokenizer.pad_token_id = tokenizer.eos_token_id

    special_tokens_to_add = []
    if model.use_start_thought_token:
        special_tokens_to_add.append("<|startthought|>")
    if model.use_end_thought_token:
        special_tokens_to_add.append("<|endthought|>")
    if special_tokens_to_add:
        tokenizer.add_special_tokens({"additional_special_tokens": special_tokens_to_add})
    model.tokenizer = tokenizer
    for name, module in model.named_modules():
        if "embed" in name:
            print(module, flush=True)

    model.gumbel_detach = gumbel_detach
    model.include_policy_loss = include_policy_loss
    model.use_end_thought_token = use_end_thought_token
    model.use_start_thought_token = use_start_thought_token
    model.n_ahead = n_ahead
    model.n_ahead_talk = n_ahead_talk
    model.n_passes = n_passes
    model.residual_think_head = residual_think_head
    model.optimize_lm_head_only_at_start = optimize_lm_head_only_at_start
    model.gumbel_temperature = gumbel_temperature
    model.original_mode = original
    model.config_params = params
    model.run_start = int(time.time())
    model.train()
    return model

max_seq_length = 8192
run_id = int(time.time())
training_args = TrainingArguments(
    output_dir="./out",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    gradient_checkpointing=False,
    gradient_accumulation_steps=1,
    optim="lion_32bit",
    logging_steps=1,
    save_strategy="steps",
    save_steps=100,
    max_steps=-1,
    # auto_find_batch_size=True,
    weight_decay=0.001,
    bf16=True,
    
    tf32=True,
    learning_rate=1e-07,
    max_grad_norm=0,
    warmup_steps=20,
    lr_scheduler_type="constant",
    push_to_hub=False,
    report_to="wandb"
)

peft_config = LoraConfig(
    r = 8, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules =["up_proj", "down_proj", "gate_proj"],
    lora_alpha = 32,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none", 
    use_dora=False,
    task_type="CAUSAL_LM"
)

torch.autograd.set_detect_anomaly(True)

class CustomSFTTrainer(SFTTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.beta = 0.9  # momentum factor
        self.clip_factor = 1.0  # clipping factor
        self.moving_avg = 0.0

    def training_step(self, model, inputs):
        model.train()
        inputs = self._prepare_inputs(inputs)

        outputs = model(**inputs)
        loss = outputs.loss if isinstance(outputs, dict) else outputs[0]

        if self.args.gradient_accumulation_steps > 1:
            loss = loss / self.args.gradient_accumulation_steps

        loss.backward()

        # Compute gradients and their norm
        grad_norm = torch.sqrt(sum(p.grad.data.norm().to(model.device)**2 for p in model.parameters() if p.grad is not None))

        # Update moving average and apply gradient clipping
        if self.state.global_step == 0:
            self.moving_avg = grad_norm
        else:
            self.moving_avg = self.beta * self.moving_avg + (1 - self.beta) * grad_norm

        if grad_norm > self.clip_factor * self.moving_avg:
            clip_coef = (self.clip_factor * self.moving_avg / grad_norm).item()
            for param in model.parameters():
                if param.grad is not None:
                    param.grad.data.mul_(clip_coef)

        if (self.state.global_step + 1) % self.args.gradient_accumulation_steps == 0:
            self.optimizer.step()
            self.lr_scheduler.step()
            model.zero_grad()
            self.state.global_step += 1

        # Return the loss as a Tensor
        return loss
    
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model_init(None)  

trainer = CustomSFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=model.tokenizer,
    max_seq_length=max_seq_length,
    peft_config=peft_config,
)

trainer.train()