diff --git a/train.py b/train.py index 931b116..6de6837 100644 --- a/train.py +++ b/train.py @@ -54,6 +54,23 @@ def train(config_path): trust_remote_code=True, low_cpu_mem_usage=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 + + print("✓ Success: QLoRA 4-bit") except Exception as e: errors.append(("QLoRA 4-bit", e)) @@ -71,6 +88,23 @@ def train(config_path): trust_remote_code=True, low_cpu_mem_usage=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 + + print("✓ Success: bf16 GPU") except Exception as e: errors.append(("bf16 GPU", e)) @@ -88,6 +122,24 @@ def train(config_path): trust_remote_code=True, low_cpu_mem_usage=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 + + + print("✓ Success: bf16 CPU") except Exception as e: errors.append(("bf16 CPU", e))