Disable training on verifier_lm_head

fynnsu · fynnsu · commit 0e33c78f9bc1 · 2025-10-03T16:10:45.000-04:00
Signed-off-by: Fynn Schmitt-Ulms &lt;fschmitt@redhat.com&gt;
diff --git a/src/speculators/train/eagle3/core.py b/src/speculators/train/eagle3/core.py
@@ -55,6 +55,7 @@ def __init__(
         self.verifier_lm_head = torch.nn.Linear(
             hidden_size, self.draft_vocab_size, bias=False
         )
+        self.verifier_lm_head.weight.requires_grad = False
         self.lm_head = torch.nn.Linear(hidden_size, self.draft_vocab_size, bias=False)
         # shape: [hidden_size, draft_vocab_size]
 

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@ def __init__(`
`55`	`55`	`self.verifier_lm_head = torch.nn.Linear(`
`56`	`56`	`hidden_size, self.draft_vocab_size, bias=False`
`57`	`57`	`)`
	`58`	`+ self.verifier_lm_head.weight.requires_grad = False`
`58`	`59`	`self.lm_head = torch.nn.Linear(hidden_size, self.draft_vocab_size, bias=False)`
`59`	`60`	`# shape: [hidden_size, draft_vocab_size]`
`60`	`61`