feat: add --check-only flag to validate setup without training

train-on-this-server.sh

@@ -78,6 +78,8 @@ echo "Press Ctrl+C to stop (model will be saved at checkpoint)."
 echo ""
 
 # Use torchrun for distributed training (2 GPUs)
+# Add --check-only to validate without training:
+# torchrun --nproc_per_node=2 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml --check-only
 torchrun --nproc_per_node=2 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml
 
 echo ""

training/scripts/train.py

@@ -122,10 +122,92 @@ def main():
     parser = argparse.ArgumentParser(description="Train LoRA adapter")
     parser.add_argument("--config", type=str, default="configs/llama2-7b-lora.yaml", 
                        help="Training configuration file")
+    parser.add_argument("--check-only", action="store_true",
+                       help="Validate config and dependencies without training")
     args = parser.parse_args()
     
+    if args.check_only:
+        check_setup(args.config)
+    else:
         train(args.config)
 
 
+def check_setup(config_path):
+    """Validate config and dependencies without loading model."""
+    print("=== Checking Setup ===")
+    
+    # Check config file
+    print(f"\n1. Config file: {config_path}")
+    if not Path(config_path).exists():
+        print(f"   ERROR: Config file not found: {config_path}")
+        return False
+    print("   ✓ Config file exists")
+    
+    # Load and validate config
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    print("   ✓ Config file is valid YAML")
+    
+    # Check model path
+    print(f"\n2. Model path: {config['base_model']}")
+    if Path(config['base_model']).exists():
+        print(f"   ✓ Model path exists: {config['base_model']}")
+    else:
+        print(f"   ⚠ Model path not found: {config['base_model']}")
+        print("   (Will download from HuggingFace during training)")
+    
+    # Check dataset files
+    print("\n3. Dataset files:")
+    repo_root = Path(__file__).parent.parent.parent
+    train_path = repo_root / "training" / "data" / "train.jsonl"
+    test_path = repo_root / "training" / "data" / "test.jsonl"
+    
+    if train_path.exists():
+        print(f"   ✓ Train data: {train_path}")
+    else:
+        print(f"   ✗ Train data missing: {train_path}")
+    
+    if test_path.exists():
+        print(f"   ✓ Test data: {test_path}")
+    else:
+        print(f"   ✗ Test data missing: {test_path}")
+    
+    # Check GPU
+    print("\n4. GPU:")
+    import torch
+    if torch.cuda.is_available():
+        print(f"   ✓ GPU available: {torch.cuda.get_device_name(0)}")
+        print(f"   ✓ VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+        print(f"   ✓ GPU count: {torch.cuda.device_count()}")
+    else:
+        print("   ✗ No GPU detected!")
+        return False
+    
+    # Check required packages
+    print("\n5. Required packages:")
+    packages = ['transformers', 'datasets', 'trl', 'peft', 'accelerate', 'deepspeed']
+    for pkg in packages:
+        try:
+            __import__(pkg)
+            print(f"   ✓ {pkg}")
+        except ImportError:
+            print(f"   ✗ {pkg} not installed")
+    
+    # Check DeepSpeed config
+    print("\n6. DeepSpeed config:")
+    if 'deepspeed_config' in config:
+        print("   ✓ DeepSpeed config present")
+        ds_config = config['deepspeed_config']
+        if 'zero_optimization' in ds_config:
+            stage = ds_config['zero_optimization'].get('stage', 'N/A')
+            print(f"   ✓ ZeRO stage: {stage}")
+    else:
+        print("   ✗ No DeepSpeed config found")
+    
+    print("\n=== Check Complete ===")
+    print("If all checks pass, you can run training with:")
+    print(f"  bash train-on-this-server.sh")
+
+
 if __name__ == "__main__":
     main()