NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-28 17:29:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

checkpointing fixes

Browse source
This commit is contained in:
Jai Suphavadeeprasit 2026-01-29 11:41:24 -05:00
parent b9414e4076
commit 04652fd97c
1 changed files with 2 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
example_trainer/model.py
View file

@ -48,6 +48,8 @@ def load_model_and_tokenizer(
if model is not None: if model is not None:
print("[Setup] ✓ Single-copy mode active - using vLLM's tensors directly!") print("[Setup] ✓ Single-copy mode active - using vLLM's tensors directly!")
# Enable gradient checkpointing to save memory (was missing before!)
_setup_gradient_checkpointing(model, config)
model.train() model.train()
return model, tokenizer return model, tokenizer
else: else: