fix FSDP2 test case failure on XPU

[yao-matrix]

bug

when run pytest -rA tests/fsdp/test_fsdp.py::FSDP2IntegrationTest::test_checkpointing on XPU, fail message as

stderr: [rank1]: main()
stderr: [rank1]: File "/opt/venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py", line 265, in main
stderr: [rank1]: training_function(config, args)
stderr: [rank1]: File "/opt/venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py", line 171, in training_function
stderr: [rank1]: accelerator.load_state(args.resume_from_checkpoint)
stderr: [rank1]: File "/opt/venv/lib/python3.12/site-packages/accelerate/accelerator.py", line 3714, in load_state
stderr: [rank1]: self.scaler._lazy_init_scale_growth_tracker(self.scaler._device)
stderr: [rank1]: File "/opt/venv/lib/python3.12/site-packages/torch/amp/grad_scaler.py", line 171, in _lazy_init_scale_growth_tracker
stderr: [rank1]: assert self._growth_tracker is None, "_growth_tracker initialized before _scale"
stderr: [rank1]: ^^^^^^^^^^^^^^^^^^^^
stderr: [rank1]: AttributeError: 'GradScaler' object has no attribute '_growth_tracker'. Did you mean: '_get_growth_tracker'?

root cause

when instantiate GradScaler in get_fsdp2_grad_scaler, the device is not passed, and will not because device is not in GradScalerKwargs either. In this case, PyTorch will defaultly set device to cuda, and lead to this issue.

fix

explicitly pass device in.

result

pass

[yao-matrix]

@SunMarc , pls help review, thx

[yao-matrix]

prior error message is misleading.

[S1ro1]

Thank you! LGTM

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.