blockblockblock's picture
Upload folder using huggingface_hub
9742816 verified
import torch
torch.backends.cuda.matmul.allow_tf32 = True
import random
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from transformers import TrainingArguments
from trl import SFTTrainer
from peft import LoraConfig
import time
random_seed = 42
torch.manual_seed(random_seed)
random.seed(random_seed)
dataset = load_dataset("HuggingFaceH4/deita-10k-v0-sft", split="train_sft")
n_ahead_talk_global = 2
n_passes_global = 2
n_ahead_global = 2
n_examples = 0
full_batch_size = 2
eval_and_logging_steps = 2
save_steps = 100
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)
gradient_accumulation_steps = params.get("gradient_accumulation_steps", global_gradient_accumulation_steps)
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.float16 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,padding=False,truncation=True)
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.resize_token_embeddings(len(tokenizer))
model.tokenizer = tokenizer
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.n_tokens_print = gradient_accumulation_steps
model.gradient_accumulation_steps = gradient_accumulation_steps
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.kill_after = 100
model.train()
return model
batch_size = full_batch_size // n_passes_global
global_gradient_accumulation_steps = full_batch_size // batch_size
run_id = int(time.time())
training_args = TrainingArguments(
output_dir="./out",
num_train_epochs=3,
per_device_train_batch_size=1,
gradient_checkpointing=False,
gradient_accumulation_steps=4,
optim="adamw_torch_fused",
logging_steps=1,
save_strategy="steps",
save_steps=300,
bf16=True,
tf32=False,
# auto_find_batch_size=True
learning_rate=2e-07,
max_grad_norm=1.0, # Gradient clipping with a maximum gradient norm of 0.3
warmup_steps=100,
lr_scheduler_type="cosine",
push_to_hub=False,
)
# peft_config = LoraConfig(
# r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
# target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
# "gate_proj", "up_proj", "down_proj",],
# lora_alpha = 16,
# lora_dropout = 0, # Supports any, but = 0 is optimized
# bias = "none", # Enable Dora method
# use_dora=True,
# )
torch.autograd.set_detect_anomaly(True)
model = model_init(None) # Initialize the model
tokenizer = model.tokenizer
trainer = SFTTrainer(
args=training_args,
train_dataset=dataset,
model=model,
# peft_config=peft_config,
tokenizer=tokenizer,
)
trainer.train()