File size: 4,900 Bytes
d507e0c 8b9bd5a d507e0c 8b9bd5a d507e0c 6524016 d507e0c 562cf21 8b9bd5a d507e0c 562cf21 d507e0c 21a96f6 d507e0c 21a96f6 d507e0c 562cf21 21a96f6 8b9bd5a d507e0c 21a96f6 d507e0c 8b9bd5a d507e0c 21a96f6 d507e0c 6524016 d507e0c 562cf21 21a96f6 562cf21 d507e0c 21a96f6 d507e0c 21a96f6 d842ce9 8b9bd5a 562cf21 21a96f6 562cf21 d507e0c 8b9bd5a d507e0c 21a96f6 d507e0c 562cf21 d507e0c 8b9bd5a d507e0c 21a96f6 d507e0c 8b9bd5a d507e0c 6524016 |
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 |
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("Crystalcareai/Self-Discover-MM-Instruct-openai", split="train_sft")
n_ahead_talk_global = 4
n_passes_global = 2
n_ahead_global = 2
n_examples = 0
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.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",
# load_in_4bit=True,
# attn_implementation="flash_attention_2",
)
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.resize_token_embeddings(len(tokenizer))
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=3,
per_device_train_batch_size=1,
gradient_checkpointing=False,
gradient_accumulation_steps=6,
optim="lion_32bit",
logging_steps=1,
save_strategy="steps",
save_steps=25,
bf16=True,
# fp16=True,
tf32=False,
# epsilson=1e-05,
# beta1=0.9,
# beta2=0.95,
# auto_find_batch_size=True
learning_rate=6e-05,
max_grad_norm=0.3, # Gradient clipping with a maximum gradient norm of 0.3
warmup_ratio=0.06,
lr_scheduler_type="cosine",
push_to_hub=False,
report_to="wandb"
)
peft_config = LoraConfig(
r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
target_modules = ["q_proj", "k_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,
max_seq_length=max_seq_length,
)
trainer.train()
|