Upload train_h100.py with huggingface_hub

train_h100.py

@@ -17,11 +17,12 @@ Optimized for H100 80GB
 """
 
 import os
+from dataclasses import dataclass
+from typing import Any, Dict, List
 from datasets import load_dataset
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
-    DataCollatorForLanguageModeling,
     Trainer,
     TrainingArguments,
 )
@@ -30,6 +31,38 @@ import torch
 import trackio
 from huggingface_hub import whoami
 
+
+@dataclass
+class DataCollatorForPreTokenized:
+    """Data collator for pre-tokenized datasets with padding."""
+    pad_token_id: int
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        # Find max length in batch
+        max_length = max(len(f["input_ids"]) for f in features)
+
+        batch = {
+            "input_ids": [],
+            "attention_mask": [],
+            "labels": [],
+        }
+
+        for feature in features:
+            input_ids = feature["input_ids"]
+            attention_mask = feature["attention_mask"]
+            labels = feature.get("labels", input_ids.copy())
+
+            # Calculate padding needed
+            padding_length = max_length - len(input_ids)
+
+            # Pad sequences (right padding)
+            batch["input_ids"].append(input_ids + [self.pad_token_id] * padding_length)
+            batch["attention_mask"].append(attention_mask + [0] * padding_length)
+            batch["labels"].append(labels + [-100] * padding_length)  # -100 is ignored by loss
+
+        # Convert to tensors
+        return {k: torch.tensor(v, dtype=torch.long) for k, v in batch.items()}
+
 # Configuration
 BASE_MODEL = "Tesslate/Synthia-S1-27b"
 OUTPUT_MODEL = "Synthia-S1-27b-tool-calling"
@@ -148,7 +181,7 @@ training_args = TrainingArguments(
     report_to="trackio",
     run_name=f"synthia-tool-calling-lora-r{LORA_R}",
     bf16=True,
-    dataloader_num_workers=4,
+    dataloader_num_workers=0,  # Avoid multiprocessing issues with custom collator
     dataloader_pin_memory=True,
     seed=42,
     remove_unused_columns=False,
@@ -156,10 +189,7 @@ training_args = TrainingArguments(
 
 # Initialize trainer
 print("\nInitializing trainer...")
-data_collator = DataCollatorForLanguageModeling(
-    tokenizer=tokenizer,
-    mlm=False,
-)
+data_collator = DataCollatorForPreTokenized(pad_token_id=tokenizer.pad_token_id)
 
 trainer = Trainer(
     model=model,