Upload train_tool_calling.py with huggingface_hub

train_tool_calling.py

@@ -0,0 +1,559 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch>=2.0.0",
+#     "transformers>=4.50.0",
+#     "datasets>=2.14.0",
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "accelerate>=0.25.0",
+#     "bitsandbytes>=0.41.0",
+#     "trackio",
+#     "huggingface_hub",
+# ]
+# ///
+"""
+LoRA Fine-tuning Script: Add Tool Calling to Synthia-S1-27b
+
+This script fine-tunes Tesslate/Synthia-S1-27b with LoRA using the
+nvidia/Nemotron-Agentic-v1 tool_calling dataset.
+
+Usage:
+    # With uv (recommended)
+    uv run train_tool_calling.py
+
+    # Or with pip
+    pip install torch transformers datasets trl peft accelerate bitsandbytes trackio
+    python train_tool_calling.py
+
+Hardware Requirements:
+    - Minimum: 1x A100 80GB or 2x A10G 24GB
+    - Recommended: 1x A100 80GB for fastest training
+"""
+
+import os
+import json
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, DataCollatorForLanguageModeling
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from trl import SFTTrainer, SFTConfig
+import torch
+import trackio
+from huggingface_hub import hf_hub_download, HfApi, create_repo
+
+# ============================================================================
+# CONFIGURATION - Modify these values as needed
+# ============================================================================
+
+# Model configuration
+BASE_MODEL = "Tesslate/Synthia-S1-27b"
+OUTPUT_MODEL = "Synthia-S1-27b-tool-calling"  # Will be pushed as Codyfederer/Synthia-S1-27b-tool-calling
+
+# Dataset configuration
+DATASET_NAME = "nvidia/Nemotron-Agentic-v1"
+DATASET_SPLIT = "tool_calling"
+MAX_SAMPLES = None  # Set to a number (e.g., 10000) to limit dataset size for testing
+
+# Training hyperparameters
+NUM_EPOCHS = 1  # 1 epoch is often sufficient for large datasets
+MAX_SEQ_LENGTH = 4096  # Adjust based on your GPU memory
+BATCH_SIZE = 1  # Per device batch size
+GRADIENT_ACCUMULATION = 16  # Effective batch size = BATCH_SIZE * GRADIENT_ACCUMULATION
+LEARNING_RATE = 2e-4
+WARMUP_RATIO = 0.03
+
+# LoRA configuration
+LORA_R = 64  # LoRA rank - higher = more capacity but more memory
+LORA_ALPHA = 128  # LoRA alpha - typically 2x rank
+LORA_DROPOUT = 0.05
+
+# Quantization (4-bit for memory efficiency)
+USE_4BIT = False  # Using BF16 LoRA on H100 for better quality
+
+# Tokenized dataset caching
+TOKENIZED_DATASET_REPO = "Codyfederer/synthia-tool-calling-tokenized"
+SAVE_TOKENIZED = True  # Save tokenized dataset to Hub for reuse
+TOKENIZED_DATASET_PRIVATE = True  # Make tokenized dataset private
+LOAD_TOKENIZED_IF_EXISTS = True  # Skip tokenization if already exists on Hub
+
+# Hub configuration
+PUSH_TO_HUB = True
+HUB_PRIVATE = False  # Set to True for private model
+
+# ============================================================================
+# TRAINING SCRIPT
+# ============================================================================
+
+def tokenize_conversation(example, tokenizer, max_length):
+    """
+    Tokenize a conversation using the model's chat template.
+    Returns input_ids, attention_mask, and labels for causal LM training.
+    """
+    messages = example["messages"]
+
+    # Apply chat template to get the full text
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=False
+    )
+
+    # Tokenize the text
+    tokenized = tokenizer(
+        text,
+        truncation=True,
+        max_length=max_length,
+        padding=False,  # We'll pad later in the data collator
+        return_tensors=None,  # Return lists, not tensors
+    )
+
+    # For causal LM, labels are the same as input_ids (shifted internally by the model)
+    tokenized["labels"] = tokenized["input_ids"].copy()
+
+    return tokenized
+
+
+def main():
+    print("=" * 60)
+    print("Tool Calling Fine-tuning for Synthia-S1-27b")
+    print("=" * 60)
+
+    # Initialize Trackio for monitoring
+    trackio.init(project="synthia-tool-calling")
+
+    # Get HF username for hub_model_id
+    from huggingface_hub import whoami
+    try:
+        username = whoami()["name"]
+        hub_model_id = f"{username}/{OUTPUT_MODEL}"
+        print(f"Will push to: {hub_model_id}")
+    except Exception as e:
+        print(f"Warning: Not logged in to HF Hub ({e})")
+        print("Model will be saved locally only. Run 'huggingface-cli login' to enable Hub push.")
+        hub_model_id = OUTPUT_MODEL
+        global PUSH_TO_HUB
+        PUSH_TO_HUB = False
+
+    # -------------------------------------------------------------------------
+    # Load Tokenizer FIRST (needed for tokenization)
+    # -------------------------------------------------------------------------
+    print(f"\nLoading tokenizer from {BASE_MODEL}...")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        BASE_MODEL,
+        trust_remote_code=True,
+        padding_side="right",
+    )
+
+    # Ensure pad token is set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    print(f"Vocab size: {len(tokenizer):,}")
+
+    # -------------------------------------------------------------------------
+    # Try to Load Pre-tokenized Dataset from Hub
+    # -------------------------------------------------------------------------
+    train_dataset = None
+    eval_dataset = None
+
+    if LOAD_TOKENIZED_IF_EXISTS:
+        print(f"\nChecking for pre-tokenized dataset: {TOKENIZED_DATASET_REPO}")
+        try:
+            from datasets import load_dataset as hf_load_dataset
+
+            # Try to load the tokenized dataset
+            tokenized_ds = hf_load_dataset(TOKENIZED_DATASET_REPO)
+
+            # Check if it has the required columns (input_ids, attention_mask)
+            if "train" in tokenized_ds and "input_ids" in tokenized_ds["train"].column_names:
+                print("  Found pre-tokenized dataset with input_ids!")
+                train_dataset = tokenized_ds["train"]
+                eval_dataset = tokenized_ds.get("test", tokenized_ds.get("validation"))
+                print(f"  Train samples: {len(train_dataset):,}")
+                if eval_dataset:
+                    print(f"  Eval samples: {len(eval_dataset):,}")
+            else:
+                print("  Dataset exists but is not tokenized (no input_ids column)")
+                print("  Will re-tokenize and save...")
+        except Exception as e:
+            print(f"  Could not load pre-tokenized dataset: {e}")
+            print("  Will tokenize from scratch...")
+
+    # -------------------------------------------------------------------------
+    # Load and Tokenize Dataset (if not loaded from Hub)
+    # -------------------------------------------------------------------------
+    if train_dataset is None:
+        print(f"\nLoading dataset: {DATASET_NAME} ({DATASET_SPLIT} split)...")
+
+        # Download the JSONL file directly from the dataset repo
+        jsonl_file = f"data/{DATASET_SPLIT}.jsonl"
+        print(f"Downloading {jsonl_file}...")
+
+        local_path = hf_hub_download(
+            repo_id=DATASET_NAME,
+            filename=jsonl_file,
+            repo_type="dataset"
+        )
+        print(f"Downloaded to: {local_path}")
+
+        # Load JSONL manually to handle schema inconsistencies
+        print("Loading and processing JSONL file...")
+        processed_examples = []
+        skipped = 0
+
+        with open(local_path, 'r', encoding='utf-8') as f:
+            for line_num, line in enumerate(f):
+                if line_num % 50000 == 0:
+                    print(f"  Processed {line_num:,} lines...")
+                try:
+                    example = json.loads(line.strip())
+                    messages = example.get("messages", [])
+
+                    # Convert messages to consistent format
+                    formatted_messages = []
+                    for msg in messages:
+                        role = msg.get("role", "user")
+                        content = msg.get("content", "")
+
+                        # Handle content that might be a list or complex object
+                        if isinstance(content, list):
+                            # For tool calls, content is often a list of dicts
+                            parts = []
+                            for item in content:
+                                if isinstance(item, dict):
+                                    if "text" in item:
+                                        parts.append(item["text"])
+                                    else:
+                                        parts.append(json.dumps(item))
+                                else:
+                                    parts.append(str(item))
+                            content = "\n".join(parts) if parts else ""
+                        elif isinstance(content, dict):
+                            content = json.dumps(content)
+                        elif content is None:
+                            content = ""
+                        else:
+                            content = str(content)
+
+                        formatted_messages.append({
+                            "role": role,
+                            "content": content
+                        })
+
+                    # Ensure proper role alternation for chat template
+                    # Merge consecutive messages with same role, handle tool messages
+                    if formatted_messages:
+                        merged_messages = []
+                        for msg in formatted_messages:
+                            role = msg["role"]
+                            content = msg["content"]
+
+                            # Map tool role to assistant (tool responses are from assistant's perspective)
+                            if role == "tool":
+                                role = "user"  # Tool output is provided to the model like user input
+                                content = f"[Tool Result]\n{content}"
+
+                            # If same role as previous, merge content
+                            if merged_messages and merged_messages[-1]["role"] == role:
+                                merged_messages[-1]["content"] += f"\n\n{content}"
+                            else:
+                                merged_messages.append({"role": role, "content": content})
+
+                        # Ensure conversation starts with user and alternates
+                        if merged_messages and merged_messages[0]["role"] != "user":
+                            # Prepend a placeholder user message if starts with assistant
+                            merged_messages.insert(0, {"role": "user", "content": "[Start]"})
+
+                        processed_examples.append({"messages": merged_messages})
+
+                except Exception as e:
+                    skipped += 1
+                    if skipped < 5:
+                        print(f"  Warning: Skipped line {line_num}: {e}")
+
+        print(f"Loaded {len(processed_examples):,} examples (skipped {skipped})")
+
+        # Create dataset from processed examples
+        dataset = Dataset.from_list(processed_examples)
+        print(f"Dataset size: {len(dataset):,} examples")
+
+        if MAX_SAMPLES and len(dataset) > MAX_SAMPLES:
+            dataset = dataset.shuffle(seed=42).select(range(MAX_SAMPLES))
+            print(f"Limited to {MAX_SAMPLES:,} samples for training")
+
+        # Create train/eval split
+        split_dataset = dataset.train_test_split(test_size=0.02, seed=42)
+        train_dataset = split_dataset["train"]
+        eval_dataset = split_dataset["test"]
+
+        print(f"Train samples: {len(train_dataset):,}")
+        print(f"Eval samples: {len(eval_dataset):,}")
+
+        # -------------------------------------------------------------------------
+        # TOKENIZE the dataset (this is the key step!)
+        # -------------------------------------------------------------------------
+        print(f"\nTokenizing dataset with max_length={MAX_SEQ_LENGTH}...")
+        print("This may take a while for large datasets...")
+
+        # Tokenize train dataset
+        train_dataset = train_dataset.map(
+            lambda x: tokenize_conversation(x, tokenizer, MAX_SEQ_LENGTH),
+            remove_columns=["messages"],  # Remove text, keep only tokens
+            num_proc=4,  # Parallelize
+            desc="Tokenizing train",
+        )
+
+        # Tokenize eval dataset
+        eval_dataset = eval_dataset.map(
+            lambda x: tokenize_conversation(x, tokenizer, MAX_SEQ_LENGTH),
+            remove_columns=["messages"],
+            num_proc=4,
+            desc="Tokenizing eval",
+        )
+
+        print(f"Tokenization complete!")
+        print(f"Train dataset columns: {train_dataset.column_names}")
+        print(f"Sample input_ids length: {len(train_dataset[0]['input_ids'])}")
+
+        # Save TOKENIZED dataset to Hub for reuse
+        if SAVE_TOKENIZED:
+            print(f"\nSaving TOKENIZED dataset to Hub: {TOKENIZED_DATASET_REPO}")
+            try:
+                # Create the repo if it doesn't exist (private!)
+                api = HfApi()
+                try:
+                    create_repo(
+                        TOKENIZED_DATASET_REPO,
+                        repo_type="dataset",
+                        private=TOKENIZED_DATASET_PRIVATE,
+                        exist_ok=True
+                    )
+                    print(f"  Created/verified repo (private={TOKENIZED_DATASET_PRIVATE})")
+
+                    # Try to update visibility if repo already exists
+                    if TOKENIZED_DATASET_PRIVATE:
+                        try:
+                            api.update_repo_visibility(
+                                TOKENIZED_DATASET_REPO,
+                                repo_type="dataset",
+                                private=True
+                            )
+                            print(f"  Ensured repo is private")
+                        except Exception:
+                            pass  # Ignore if already private or no permission
+                except Exception as e:
+                    print(f"  Repo creation note: {e}")
+
+                # Reset format to ensure data is serializable (not torch tensors)
+                train_dataset.reset_format()
+                eval_dataset.reset_format()
+
+                # Verify the data looks correct before pushing
+                print(f"  Verifying tokenized data...")
+                print(f"    Train columns: {train_dataset.column_names}")
+                print(f"    Sample input_ids type: {type(train_dataset[0]['input_ids'])}")
+                print(f"    Sample input_ids length: {len(train_dataset[0]['input_ids'])}")
+                print(f"    First 10 tokens: {train_dataset[0]['input_ids'][:10]}")
+
+                # Push tokenized datasets to Hub (private is set at repo creation)
+                print(f"  Pushing train split ({len(train_dataset):,} examples)...")
+                train_dataset.push_to_hub(
+                    TOKENIZED_DATASET_REPO,
+                    split="train",
+                )
+                print(f"  Pushing test split ({len(eval_dataset):,} examples)...")
+                eval_dataset.push_to_hub(
+                    TOKENIZED_DATASET_REPO,
+                    split="test",
+                )
+                print(f"  SUCCESS! Saved TOKENIZED data to: https://huggingface.co/datasets/{TOKENIZED_DATASET_REPO}")
+                print(f"  Columns saved: {train_dataset.column_names}")
+                print(f"  Dataset is private: {TOKENIZED_DATASET_PRIVATE}")
+
+                # Verify the upload by trying to load it back
+                print(f"  Verifying upload...")
+                try:
+                    from datasets import load_dataset as verify_load
+                    verify_ds = verify_load(TOKENIZED_DATASET_REPO, split="train", streaming=True)
+                    sample = next(iter(verify_ds))
+                    if "input_ids" in sample:
+                        print(f"  VERIFIED: Dataset contains input_ids with {len(sample['input_ids'])} tokens")
+                    else:
+                        print(f"  WARNING: Dataset uploaded but input_ids not found in columns: {list(sample.keys())}")
+                except Exception as ve:
+                    print(f"  Could not verify upload: {ve}")
+
+            except Exception as e:
+                print(f"  ERROR saving to Hub: {e}")
+                import traceback
+                traceback.print_exc()
+                print("  Continuing with training anyway...")
+
+    # -------------------------------------------------------------------------
+    # Load Model with Quantization
+    # -------------------------------------------------------------------------
+    print(f"\nLoading model: {BASE_MODEL}...")
+
+    if USE_4BIT:
+        print("Using 4-bit quantization (QLoRA)")
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+        )
+    else:
+        bnb_config = None
+
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",  # Use PyTorch's Scaled Dot Product Attention
+    )
+
+    if USE_4BIT:
+        model = prepare_model_for_kbit_training(model)
+
+    print(f"Model loaded. Parameters: {model.num_parameters():,}")
+
+    # -------------------------------------------------------------------------
+    # Configure LoRA
+    # -------------------------------------------------------------------------
+    print(f"\nConfiguring LoRA (r={LORA_R}, alpha={LORA_ALPHA})...")
+
+    # Target modules for Gemma 3 architecture
+    target_modules = [
+        "q_proj", "k_proj", "v_proj", "o_proj",  # Attention
+        "gate_proj", "up_proj", "down_proj",      # MLP
+    ]
+
+    lora_config = LoraConfig(
+        r=LORA_R,
+        lora_alpha=LORA_ALPHA,
+        lora_dropout=LORA_DROPOUT,
+        target_modules=target_modules,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    # -------------------------------------------------------------------------
+    # Training Configuration
+    # -------------------------------------------------------------------------
+    print("\nConfiguring training...")
+
+    training_args = SFTConfig(
+        output_dir=f"./{OUTPUT_MODEL}",
+
+        # Training params
+        num_train_epochs=NUM_EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE,
+        per_device_eval_batch_size=BATCH_SIZE,
+        gradient_accumulation_steps=GRADIENT_ACCUMULATION,
+
+        # Optimizer
+        learning_rate=LEARNING_RATE,
+        lr_scheduler_type="cosine",
+        warmup_ratio=WARMUP_RATIO,
+        weight_decay=0.01,
+        optim="adamw_torch",
+
+        # Memory optimization
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        max_grad_norm=1.0,
+
+        # Sequence length
+        max_length=MAX_SEQ_LENGTH,
+        packing=False,  # Disable packing for tool calling (preserve conversation structure)
+
+        # Evaluation
+        eval_strategy="steps",
+        eval_steps=500,
+
+        # Saving
+        save_strategy="steps",
+        save_steps=500,
+        save_total_limit=3,
+
+        # Hub
+        push_to_hub=PUSH_TO_HUB,
+        hub_model_id=hub_model_id if PUSH_TO_HUB else None,
+        hub_strategy="checkpoint",
+        hub_private_repo=HUB_PRIVATE,
+
+        # Logging
+        logging_steps=10,
+        report_to="trackio",
+        run_name=f"lora-r{LORA_R}-lr{LEARNING_RATE}",
+
+        # Performance
+        bf16=True,
+        dataloader_num_workers=4,
+        dataloader_pin_memory=True,
+
+        # Reproducibility
+        seed=42,
+    )
+
+    # -------------------------------------------------------------------------
+    # Initialize Trainer
+    # -------------------------------------------------------------------------
+    print("\nInitializing trainer...")
+
+    # Create data collator for padding pre-tokenized data
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False,  # Causal LM, not masked LM
+    )
+
+    # Check if dataset is pre-tokenized
+    is_pretokenized = "input_ids" in train_dataset.column_names
+    print(f"Dataset is pre-tokenized: {is_pretokenized}")
+    print(f"Dataset columns: {train_dataset.column_names}")
+
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        processing_class=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # -------------------------------------------------------------------------
+    # Train!
+    # -------------------------------------------------------------------------
+    print("\n" + "=" * 60)
+    print("Starting training...")
+    print("=" * 60 + "\n")
+
+    trainer.train()
+
+    # -------------------------------------------------------------------------
+    # Save Final Model
+    # -------------------------------------------------------------------------
+    print("\nSaving final model...")
+    trainer.save_model()
+
+    if PUSH_TO_HUB:
+        print(f"Pushing to Hub: {hub_model_id}")
+        trainer.push_to_hub()
+        print(f"\n✅ Model available at: https://huggingface.co/{hub_model_id}")
+    else:
+        print(f"Model saved locally to: ./{OUTPUT_MODEL}")
+
+    print("\n" + "=" * 60)
+    print("Training complete!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()