p0c: pass attn_implementation via config (transformers 5.0 break)

probes/p0c_fp16_ddp_probe.py

@@ -55,8 +55,14 @@ CFG = LlamaConfig(
 
 
 def make_model(attn="sdpa"):
-    m = LlamaForCausalLM(CFG, attn_implementation=attn)
-    return m.cuda()
+    # transformers 5.0 removed `attn_implementation` from PreTrainedModel.__init__ -- passing it
+    # there raises TypeError (this is what killed probe C's first attempt). It is a config field.
+    cfg = LlamaConfig(
+        hidden_size=H, num_hidden_layers=NL, num_attention_heads=NH, num_key_value_heads=NKV,
+        intermediate_size=FF, vocab_size=VOCAB, max_position_embeddings=SEQ,
+        tie_word_embeddings=True, attn_implementation=attn,
+    )
+    return LlamaForCausalLM(cfg).cuda()
 
 
 def n_params(m):
@@ -211,9 +217,10 @@ torch.cuda.set_device(rank)
 H, NL, NH, NKV, FF, VOCAB, SEQ = 768, 12, 12, 3, 2048, 50257, 1024
 cfg = LlamaConfig(hidden_size=H, num_hidden_layers=NL, num_attention_heads=NH,
                   num_key_value_heads=NKV, intermediate_size=FF, vocab_size=VOCAB,
-                  max_position_embeddings=SEQ, tie_word_embeddings=True)
+                  max_position_embeddings=SEQ, tie_word_embeddings=True,
+                  attn_implementation="sdpa")
 torch.manual_seed(1234)
-model = LlamaForCausalLM(cfg, attn_implementation="sdpa").cuda().to(torch.float32)
+model = LlamaForCausalLM(cfg).cuda().to(torch.float32)
 model = DDP(model, device_ids=[rank])
 opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
 scaler = torch.amp.GradScaler("cuda")