[VLM] feat: true on policy for vlm + fsdp (#14636)

Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
Nan Jiang
2026-01-01 16:54:39 -08:00
committed by GitHub
co-authored by Xinyuan Tong
parent 00562ee14a
commit 7254986342
12 changed files with 203 additions and 57 deletions
+2
View File
@@ -506,6 +506,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor],
captured_last_layer_outputs: Optional[List[torch.Tensor]] = None,
**kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
hidden_states, residual = (
@@ -514,6 +515,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
residual,
forward_batch,
captured_last_layer_outputs=captured_last_layer_outputs,
**kwargs,
)
)