System Info
Python : 3.10.12
Accelerate : 0.33.0
System : Pop Os
Information
Tasks
Reproduction
Currently when you do accelerator.prepare(model, optimizer) and use msamp. The below logic is applied
if self.mixed_precision == "fp8" and self.fp8_recipe_handler.backend == "MSAMP":
args = self._prepare_msamp(*args)
# MS-AMP will handle the device placement
device_placement = [False for _ in args]
result = tuple(
self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
)
result = tuple(self._prepare_one(obj, device_placement=d) for obj, d in zip(result, device_placement))
If you see device placements in always False and we assume msamp will do device placements. But that's not the case.
If you see the logic of initialize in msamp. Then don't move it to device, they just cast the layers and replace some.
If you see the example in their official repo device placement is done first and the msamp.initialize happens.
I tried running an example with accelerate it keeps the model, dataloader on cpu. But without msamp enabled behaves correctly.
Fix is to call _prepare_one method with device placement to True.
Expected behavior
Every argument passed to prepare should be placed on proper device.
PS: I would love to open a PR to fix this issue ❤️
System Info
Information
Tasks
no_trainerscript in theexamplesfolder of thetransformersrepo (such asrun_no_trainer_glue.py)Reproduction
Currently when you do accelerator.prepare(model, optimizer) and use msamp. The below logic is applied
If you see device placements in always False and we assume msamp will do device placements. But that's not the case.
If you see the logic of initialize in msamp. Then don't move it to device, they just cast the layers and replace some.
If you see the example in their official repo device placement is done first and the msamp.initialize happens.
I tried running an example with accelerate it keeps the model, dataloader on cpu. But without msamp enabled behaves correctly.
Fix is to call _prepare_one method with device placement to True.
Expected behavior
Every argument passed to prepare should be placed on proper device.
PS: I would love to open a PR to fix this issue ❤️