transformers
1e2900df - Revert cleanup approach, restore forward-pass fix for embed_tokens

Commit
4 days ago
Revert cleanup approach, restore forward-pass fix for embed_tokens The GPU cannot fit Gemma4-E2B entirely with device_map="auto", so accelerate offloads embed_tokens to CPU on every model load — not just under accumulated memory pressure. The cleanup/tearDown approach doesn't help; it just forces model reloads that hit the same issue. The correct fix is to use embed_tokens.forward() instead of direct .weight access: the forward pass triggers accelerate's AlignDevicesHook which transiently loads weights from CPU to GPU before the computation. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Author
Parents
Loading