show layer class names
This commit is contained in:
21
train.py
21
train.py
@@ -36,6 +36,27 @@ def train(config_path):
|
||||
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)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user