show layer class names

This commit is contained in:
Christian Medina
2026-07-01 23:12:28 -04:00
parent e64515e3ca
commit ce575ee71b

View File

@@ -35,6 +35,27 @@ def train(config_path):
# Load model with multiple strategies # Load model with multiple strategies
print(f"\n[INFO] Loading {config['base_model']}...") print(f"\n[INFO] Loading {config['base_model']}...")
errors = [] errors = []
model = AutoModelForCausalLM.from_pretrained(
"deepreinforce-ai/Ornith-1.0-35B",
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) # Strategy 1: QLoRA with FSDP (preferred)