Refactor train.py to utilize a comprehensive configuration structure from config.yaml, enhancing model loading, dataset handling, and trainer setup. This update centralizes parameters for model, PEFT, dataset, and training settings, improving maintainability and flexibility.

conf/config.yaml

@@ -1,6 +1,72 @@
 defaults:
   - _self_
 
-model_name: "unsloth/SmolLM2-135M-Instruct-bnb-4bit"
+# Model configuration
+model:
+  name: "unsloth/SmolLM2-135M-Instruct-bnb-4bit"
+  max_seq_length: 2048  # Auto supports RoPE Scaling internally
+  dtype: null  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+  load_in_4bit: true  # Use 4bit quantization to reduce memory usage
+
+# PEFT configuration
+peft:
+  r: 64
+  lora_alpha: 128
+  lora_dropout: 0.05
+  bias: "none"
+  use_gradient_checkpointing: "unsloth"
+  random_state: 3407
+  use_rslora: true
+  loftq_config: null
+  target_modules:
+    - "q_proj"
+    - "k_proj"
+    - "v_proj"
+    - "o_proj"
+    - "gate_proj"
+    - "up_proj"
+    - "down_proj"
+
+# Dataset configuration
+dataset:
+  validation_split: 0.1  # 10% of data for validation
+  seed: 3407  # Random seed for dataset splitting
+
+# Training configuration
+training:
+  args:
+    per_device_train_batch_size: 2
+    per_device_eval_batch_size: 2
+    gradient_accumulation_steps: 16
+    warmup_steps: 100
+    max_steps: 120
+    learning_rate: 5e-5
+    logging_steps: 1
+    save_strategy: "steps"
+    save_steps: 30
+    eval_strategy: "steps"
+    eval_steps: 30
+    save_total_limit: 2
+    optim: "adamw_8bit"
+    weight_decay: 0.01
+    lr_scheduler_type: "cosine_with_restarts"
+    seed: 3407
+    output_dir: "outputs"
+    gradient_checkpointing: true
+    load_best_model_at_end: true
+    metric_for_best_model: "eval_loss"
+    greater_is_better: false
+
+  sft:
+    dataset_num_proc: 2
+    packing: false
+    data_collator:
+      mlm: false
+      pad_to_multiple_of: 8
+
+# Output configuration
+output:
+  dir: "final_model"
+
+# Training control
 train: false
-output_dir: "final_model"

train.py

@@ -43,13 +43,6 @@ from transformers import (
 )
 from trl import SFTTrainer
 
-# Configuration
-dtype = None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
-load_in_4bit = True  # Use 4bit quantization to reduce memory usage
-max_seq_length = 2048  # Auto supports RoPE Scaling internally
-validation_split = 0.1  # 10% of data for validation
-
-
 # Setup logging
 def setup_logging():
     """Configure logging for the training process."""
@@ -90,38 +83,30 @@ def install_dependencies():
         raise
 
 
-def load_model(model_name: str) -> tuple[FastLanguageModel, AutoTokenizer]:
+def load_model(cfg: DictConfig) -> tuple[FastLanguageModel, AutoTokenizer]:
     """Load and configure the model."""
     logger.info("Loading model and tokenizer...")
     try:
         model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=model_name,
-            max_seq_length=max_seq_length,
-            dtype=dtype,
-            load_in_4bit=load_in_4bit,
+            model_name=cfg.model.name,
+            max_seq_length=cfg.model.max_seq_length,
+            dtype=cfg.model.dtype,
+            load_in_4bit=cfg.model.load_in_4bit,
         )
         logger.info("Base model loaded successfully")
 
         # Configure LoRA
         model = FastLanguageModel.get_peft_model(
             model,
-            r=64,
-            target_modules=[
-                "q_proj",
-                "k_proj",
-                "v_proj",
-                "o_proj",
-                "gate_proj",
-                "up_proj",
-                "down_proj",
-            ],
-            lora_alpha=128,
-            lora_dropout=0.05,
-            bias="none",
-            use_gradient_checkpointing="unsloth",
-            random_state=3407,
-            use_rslora=True,
-            loftq_config=None,
+            r=cfg.peft.r,
+            target_modules=cfg.peft.target_modules,
+            lora_alpha=cfg.peft.lora_alpha,
+            lora_dropout=cfg.peft.lora_dropout,
+            bias=cfg.peft.bias,
+            use_gradient_checkpointing=cfg.peft.use_gradient_checkpointing,
+            random_state=cfg.peft.random_state,
+            use_rslora=cfg.peft.use_rslora,
+            loftq_config=cfg.peft.loftq_config,
         )
         logger.info("LoRA configuration applied successfully")
 
@@ -133,6 +118,7 @@ def load_model(model_name: str) -> tuple[FastLanguageModel, AutoTokenizer]:
 
 def load_and_format_dataset(
     tokenizer: AutoTokenizer,
+    cfg: DictConfig,
 ) -> tuple[
     Union[DatasetDict, Dataset, IterableDatasetDict, IterableDataset], AutoTokenizer
 ]:
@@ -144,7 +130,7 @@ def load_and_format_dataset(
         logger.info(f"Dataset loaded successfully. Size: {len(dataset)} examples")
 
         # Split into train and validation sets
-        dataset = dataset.train_test_split(test_size=validation_split, seed=3407)
+        dataset = dataset.train_test_split(test_size=cfg.dataset.validation_split, seed=cfg.dataset.seed)
         logger.info(
             f"Dataset split into train ({len(dataset['train'])} examples) and validation ({len(dataset['test'])} examples) sets"
         )
@@ -194,47 +180,34 @@ def create_trainer(
     model: FastLanguageModel,
     tokenizer: AutoTokenizer,
     dataset: Union[DatasetDict, Dataset, IterableDatasetDict, IterableDataset],
+    cfg: DictConfig,
 ) -> Trainer:
     """Create and configure the SFTTrainer."""
     logger.info("Creating trainer...")
     try:
+        # Create TrainingArguments from config
+        training_args_dict = OmegaConf.to_container(cfg.training.args, resolve=True)
+        # Add dynamic precision settings
+        training_args_dict.update({
+            "fp16": not is_bfloat16_supported(),
+            "bf16": is_bfloat16_supported(),
+        })
+        training_args = TrainingArguments(**training_args_dict)
+
+        # Create data collator from config
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            **cfg.training.sft.data_collator,
+        )
+
         trainer = SFTTrainer(
             model=model,
             tokenizer=tokenizer,
             train_dataset=dataset["train"],
             eval_dataset=dataset["validation"],
-            dataset_num_proc=2,
-            packing=False,
-            args=TrainingArguments(
-                per_device_train_batch_size=2,
-                per_device_eval_batch_size=2,
-                gradient_accumulation_steps=16,
-                warmup_steps=100,
-                max_steps=120,
-                learning_rate=5e-5,
-                fp16=not is_bfloat16_supported(),
-                bf16=is_bfloat16_supported(),
-                logging_steps=1,
-                save_strategy="steps",
-                save_steps=30,
-                eval_strategy="steps",
-                eval_steps=30,
-                save_total_limit=2,
-                optim="adamw_8bit",
-                weight_decay=0.01,
-                lr_scheduler_type="cosine_with_restarts",
-                seed=3407,
-                output_dir="outputs",
-                gradient_checkpointing=True,
-                load_best_model_at_end=True,
-                metric_for_best_model="eval_loss",
-                greater_is_better=False,
-            ),
-            data_collator=DataCollatorForLanguageModeling(
-                tokenizer=tokenizer,
-                mlm=False,
-                pad_to_multiple_of=8,
-            ),
+            args=training_args,
+            data_collator=data_collator,
+            **cfg.training.sft,
         )
         logger.info("Trainer created successfully")
         return trainer
@@ -254,13 +227,13 @@ def main(cfg: DictConfig) -> None:
         install_dependencies()
 
         # Load model and tokenizer
-        model, tokenizer = load_model(cfg.model_name)
+        model, tokenizer = load_model(cfg)
 
         # Load and prepare dataset
-        dataset, tokenizer = load_and_format_dataset(tokenizer)
+        dataset, tokenizer = load_and_format_dataset(tokenizer, cfg)
 
         # Create trainer
-        trainer: Trainer = create_trainer(model, tokenizer, dataset)
+        trainer: Trainer = create_trainer(model, tokenizer, dataset, cfg)
 
         # Train if requested
         if cfg.train:
@@ -268,8 +241,8 @@ def main(cfg: DictConfig) -> None:
             trainer.train()
 
             # Save model
-            logger.info(f"Saving final model to {cfg.output_dir}...")
-            trainer.save_model(cfg.output_dir)
+            logger.info(f"Saving final model to {cfg.output.dir}...")
+            trainer.save_model(cfg.output.dir)
 
             # Print final metrics
             final_metrics = trainer.state.log_history[-1]