diff --git a/train-on-this-server.sh b/train-on-this-server.sh index 8c9a87f..4a6b199 100755 --- a/train-on-this-server.sh +++ b/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 "" diff --git a/training/scripts/train.py b/training/scripts/train.py index fd56ccb..c0b122a 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -122,9 +122,91 @@ 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() - train(args.config) + 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__":