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(1000)) n_ahead_talk_global = 4 n_passes_global = 1 n_ahead_global = 4 # 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 = 1024 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=8, optim="adamw_torch_fused", 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=2e-10, max_grad_norm=0, warmup_steps=20, lr_scheduler_type="cosine", push_to_hub=False, report_to="wandb" ) peft_config = LoraConfig( r = 8, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 target_modules =["q_proj", "v_proj"], lora_alpha = 32, lora_dropout = 0, # Supports any, but = 0 is optimized bias = "none", use_dora=True, 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 = SFTTrainer( model=model, args=training_args, train_dataset=dataset, tokenizer=model.tokenizer, max_seq_length=max_seq_length, peft_config=peft_config, ) trainer.train()