Skip to content

Commit 2669f3e

Browse files
committed
fix: remove a empty_cache volcengine#575 missed
1 parent e7c40b3 commit 2669f3e

File tree

1 file changed

+1
-3
lines changed

1 file changed

+1
-3
lines changed

verl/utils/fsdp_utils.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ def lambda_policy_fn(module):
112112

113113

114114
@torch.no_grad()
115-
def offload_fsdp_model_to_cpu(model: FSDP, empty_cache: bool = True):
115+
def offload_fsdp_model_to_cpu(model: FSDP):
116116
assert isinstance(model, FSDP)
117117
# lazy init FSDP model
118118
_lazy_init(model, model)
@@ -128,8 +128,6 @@ def offload_fsdp_model_to_cpu(model: FSDP, empty_cache: bool = True):
128128
# the following still keeps id(._local_shard) != id(.data)
129129
flat_param._local_shard = flat_param.data
130130
assert id(flat_param._local_shard) != id(flat_param.data)
131-
if empty_cache:
132-
torch.cuda.empty_cache()
133131

134132

135133
@torch.no_grad()

0 commit comments

Comments
 (0)