Skip to content

Commit 2053ee9

Browse files
authored
Fix CPU Offloading (#1159)
SUMMARY: - When updating the training args, `place_model_on_device` was missed and as a result, when creating the trainer (which we really should not be doing during oneshot...) the default value is left as True and the trainer tries to move the model to a gpu, if it is available. - We want this argument to be False as we handle the device map and model initialization based on the calibration needs TEST PLAN: - `cpu_offloading_fp8.py` ran to completion without issue - `mult_gpus_int8_device_map` made it past the error and is running
1 parent c2db397 commit 2053ee9

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

src/llmcompressor/args/training_arguments.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,7 @@ class TrainingArguments(HFTrainingArgs):
3030
"checkpoints will be written."
3131
},
3232
)
33+
34+
@property
35+
def place_model_on_device(self):
36+
return False

0 commit comments

Comments
 (0)