From 9dd6ab01bf1caa468ee54cd88703544caf0b4093 Mon Sep 17 00:00:00 2001 From: Will Constable Date: Fri, 9 Feb 2024 16:50:10 -0800 Subject: [PATCH] Move to cuda unconditionally so pp-only run works (#50) --- torchtrain/parallelisms/parallelize_llama.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torchtrain/parallelisms/parallelize_llama.py b/torchtrain/parallelisms/parallelize_llama.py index d6db313f..399533e2 100644 --- a/torchtrain/parallelisms/parallelize_llama.py +++ b/torchtrain/parallelisms/parallelize_llama.py @@ -185,4 +185,6 @@ def parallelize_llama(model, world_mesh, parallel_dims, args): rank0_log("Applied FSDP to the model...") + # redundant if FSDP is enabled, but ensure the model is on device regardless of which parallelisms were used + model.cuda() return model