Upload train_glm47_flash.py with huggingface_hub

train_glm47_flash.py

@@ -0,0 +1,210 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch>=2.0.0",
+#     "transformers>=4.45.0",
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "accelerate>=0.24.0",
+#     "datasets",
+#     "trackio",
+#     "bitsandbytes",
+# ]
+# ///
+
+"""
+Fine-tune GLM-4.7-Flash on Unblinded Mastery dataset for QA and instruction following.
+Using TRL SFTTrainer with LoRA on A100-80GB.
+"""
+
+import torch
+import trackio
+from datasets import load_dataset
+from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from trl import SFTTrainer, SFTConfig
+
+# Configuration
+MODEL_NAME = "zai-org/GLM-4.7-Flash"
+DATASET_NAME = "LordNeel/unblinded-mastery-sharegpt"
+OUTPUT_MODEL = "LordNeel/GLM-4.7-Flash-Unblinded-Mastery"
+
+print("=" * 60)
+print("GLM-4.7-Flash Fine-tuning for Unblinded Mastery")
+print("=" * 60)
+
+# Load dataset
+print("\nLoading dataset...")
+dataset = load_dataset(DATASET_NAME, split="train")
+print(f"Dataset loaded: {len(dataset)} examples")
+
+# Create train/eval split
+print("Creating train/eval split...")
+dataset_split = dataset.train_test_split(test_size=0.05, seed=42)
+train_dataset = dataset_split["train"]
+eval_dataset = dataset_split["test"]
+print(f"   Train: {len(train_dataset)} examples")
+print(f"   Eval: {len(eval_dataset)} examples")
+
+# 4-bit quantization config for memory efficiency
+print("\nSetting up 4-bit quantization...")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+# Load tokenizer
+print("\nLoading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+print(f"Tokenizer loaded. Vocab size: {len(tokenizer)}")
+
+# Load model with 4-bit quantization
+print("\nLoading model with 4-bit quantization...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+print("Model loaded!")
+
+# Prepare model for k-bit training
+model = prepare_model_for_kbit_training(model)
+
+# Find all linear layer names for LoRA
+print("\nFinding linear layers for LoRA...")
+def find_all_linear_names(model):
+    cls = torch.nn.Linear
+    lora_module_names = set()
+    for name, module in model.named_modules():
+        if isinstance(module, cls):
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+    # Remove output layer
+    if 'lm_head' in lora_module_names:
+        lora_module_names.remove('lm_head')
+    return list(lora_module_names)
+
+target_modules = find_all_linear_names(model)
+print(f"   Found target modules: {target_modules}")
+
+# LoRA configuration
+print("\nConfiguring LoRA...")
+peft_config = LoraConfig(
+    r=32,
+    lora_alpha=64,
+    lora_dropout=0.05,
+    bias="none",
+    task_type=TaskType.CAUSAL_LM,
+    target_modules=target_modules,
+)
+
+# Apply LoRA
+model = get_peft_model(model, peft_config)
+model.print_trainable_parameters()
+
+# Format function for ShareGPT conversations
+def format_sharegpt(example):
+    """Format ShareGPT conversations to chat template."""
+    messages = []
+    for turn in example["conversations"]:
+        role_map = {"system": "system", "human": "user", "gpt": "assistant"}
+        role = role_map.get(turn["from"], turn["from"])
+        messages.append({"role": role, "content": turn["value"]})
+
+    # Apply chat template
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+    return {"text": text}
+
+# Format datasets
+print("\nFormatting datasets...")
+train_dataset = train_dataset.map(format_sharegpt, remove_columns=train_dataset.column_names)
+eval_dataset = eval_dataset.map(format_sharegpt, remove_columns=eval_dataset.column_names)
+print("Datasets formatted!")
+
+# Training configuration
+print("\nConfiguring training...")
+training_config = SFTConfig(
+    # Hub settings - CRITICAL for saving
+    output_dir=OUTPUT_MODEL.split("/")[-1],
+    push_to_hub=True,
+    hub_model_id=OUTPUT_MODEL,
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # Training parameters
+    num_train_epochs=3,
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=2,
+    gradient_accumulation_steps=8,  # Effective batch size: 16
+    learning_rate=2e-4,
+    max_seq_length=2048,
+
+    # Memory optimization
+    gradient_checkpointing=True,
+    gradient_checkpointing_kwargs={"use_reentrant": False},
+
+    # Logging & checkpointing
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=100,
+    save_total_limit=3,
+
+    # Evaluation
+    eval_strategy="steps",
+    eval_steps=100,
+
+    # Optimization
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    optim="paged_adamw_8bit",
+
+    # Precision
+    bf16=True,
+    fp16=False,
+
+    # Monitoring
+    report_to="trackio",
+    project="unblinded-mastery-finetuning",
+    run_name="glm47flash-sft-lora",
+
+    # Dataset
+    dataset_text_field="text",
+    packing=False,
+)
+
+# Initialize trainer
+print("\nInitializing trainer...")
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    args=training_config,
+    tokenizer=tokenizer,
+    peft_config=None,  # Already applied above
+)
+
+# Train
+print("\n" + "=" * 60)
+print("STARTING TRAINING")
+print("=" * 60)
+trainer.train()
+
+# Save and push to hub
+print("\nSaving model to Hub...")
+trainer.save_model()
+trainer.push_to_hub()
+
+# Finish tracking
+trackio.finish()
+
+print("\n" + "=" * 60)
+print("TRAINING COMPLETE!")
+print(f"Model saved to: https://huggingface.co/{OUTPUT_MODEL}")
+print(f"View metrics at: https://huggingface.co/spaces/LordNeel/trackio")
+print("=" * 60)