[Spec] Wire DFLASH aux-hidden capture into the Qwen3.5 text-only wrapper (#34771)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
6f005e4da1
commit
a5ba081fbb
@@ -118,6 +118,16 @@ class Qwen3_5ForCausalLM(nn.Module):
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):
|
||||
if not self.pp_group.is_last_rank:
|
||||
return
|
||||
if layers_to_capture is None:
|
||||
raise ValueError(
|
||||
"DFLASH requires explicit layer ids for aux hidden capture."
|
||||
)
|
||||
self.capture_aux_hidden_states = True
|
||||
self.model.set_dflash_layers_to_capture(layers_to_capture)
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user