Create train.py

train.py

@@ -0,0 +1,142 @@
+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 = 4
+n_passes_global = 2
+n_ahead_global = 12
+n_examples = 1_000
+full_batch_size = 8
+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 = "Crystalcareai/Quiet-Star-Custom"
+    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,
+        load_in_4bit=True,     
+    )
+    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,
+    optim="adamw_bnb_8bit",
+    logging_steps=2,
+    save_strategy="steps",
+    save_steps=300,
+
+    bf16=True,
+    tf32=True,
+    learning_rate=2e-4,
+    max_grad_norm=0.3,
+    warmup_ratio=0.00,
+    lr_scheduler_type="constant",
+    push_to_hub=False,
+)
+
+peft_config = LoraConfig(
+        lora_alpha=16,
+        lora_dropout=0.05,
+        r=32,
+        bias="none",
+        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj",],
+        task_type="CAUSAL_LM",
+        use_dora=False,  # Enable Dora method
+)
+
+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()