Skip to content

Commit 2b7a3a1

Browse files
committed
fixed mcore offload
Signed-off-by: xhx1022 <1737006628@qq.com>
1 parent e3b187a commit 2b7a3a1

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

verl/workers/engine/megatron/transformer_impl.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,8 @@ def initialize(self):
319319
if self.engine_config.forward_only:
320320
self.optimizer = None
321321
self.lr_scheduler = None
322+
self.to(device="cpu", model=self._is_offload_param, optimizer=False, grad=False)
323+
log_gpu_memory_usage("After offload model during init (forward_only)", logger=logger)
322324
return
323325

324326
self.optimizer = self._build_optimizer()

0 commit comments

Comments
 (0)