File size: 6,561 Bytes
b0e4fff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
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() |