Skip to content

Commit

Permalink
Move to cuda unconditionally so pp-only run works
Browse files Browse the repository at this point in the history
  • Loading branch information
wconstab committed Feb 9, 2024
1 parent 7a8a9ec commit 449c824
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions torchtrain/parallelisms/parallelize_llama.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,9 @@ def parallelize_llama(model, world_mesh, parallel_dims, args):
# wrap the rest layers with FSDP
model = wrap(model.cuda())

# redundant if FSDP is used, but ensure the model is on device consistently regardless with parallelisms were used
model.cuda()

rank0_log("Applied parallelisms to the model...")

return model

0 comments on commit 449c824

Please sign in to comment.