diff --git a/train.py b/train.py index 931b116..7f2b876 100644 --- a/train.py +++ b/train.py @@ -35,6 +35,27 @@ def train(config_path): # Load model with multiple strategies print(f"\n[INFO] Loading {config['base_model']}...") errors = [] + + + model = AutoModelForCausalLM.from_pretrained( + "YOUR_MODEL_PATH", + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ) + print(type(model)) + + classes = set() + + for _, module in model.named_modules(): + classes.add(module.__class__.__name__) + + print("\n=== Transformer classes ===") + for cls in sorted(classes): + if "Block" in cls or "Layer" in cls or "Decoder" in cls: + print(cls) + + raise SystemExit + # ------------------------------------------------------------------ # Strategy 1: QLoRA with FSDP (preferred)