[release/2.14] Muon optimizer state_dict cross-device numeric tolerance failure
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 13h
- Merged PRs (30d)
- 107
Description
### 🐛 Describe the bug
Cases:
~~op_ut,third_party.torch-xpu-ops.test.xpu.test_optim_xpu.TestOptimRenewedXPU,test_state_dict_cross_device_Muon_xpu_float32~~
## Affected Test Cases
- `test_state_dict_cross_device_Muon_xpu_float32` (`TestOptimRenewedXPU` in `test_optim_xpu.py`)
## Root Cause
The Muon optimizer's `test_state_dict_cross_device` test saves optimizer state on one device and loads it on another, then compares the resulting parameter tensors. The comparison fails with a small numeric tolerance violation: greatest absolute difference of `1.94e-05` (allowed: `1e-05`) and greatest relative difference of `1.88e-05` (allowed: `1.3e-06`). This is a precision issue specific to the Muon optimizer's cross-device state restoration on XPU. The Muon optimizer (which applies Newton-Schulz orthogonalization) accumulates floating-point differences when state tensors are transferred between CPU and XPU, likely due to differences in intermediate computation precision during the orthogonalization steps. This may require either a tolerance override for the XPU Muon cross-device test or investigation into the precision of the Newton-Schulz iteration on XPU.
## Reproducer
```bash
source ~/.bashrc && cd third_party/torch-xpu-ops/test/xpu && python -m pytest -sxv test_optim_xpu.py -k "test_state_dict_cross_device_Muon_xpu_float32"
```
## Representative Error
```
Traceback (most recent call last):
File "/home/gta/Projects/pytorch/third_party/torch-xpu-ops/test/xpu/../../../../test/test_optim.py", line 1887, in test_state_dict_cross_device
self.assertEqual(params, params_device)
File "/home/gta/Projects/pytorch/torch/testing/_internal/common_utils.py", line 4608, in assertEqual
raise error_metas.pop()[0].to_error(
AssertionError: Tensor-likes are not close!
Mismatched elements: 1 / 6 (16.7%)
Greatest absolute difference: 1.9431114196777344e-05 at index (1, 0) (up to 1e-05 allowed)
Greatest relative difference: 1.8793216440826654e-05 at index (1, 0) (up to 1.3e-06 allowed)
The failure occurred for item [0]
```
### Versions
- PyTorch version: 2.14.0a0+git4dbabff
- Python: 3.10.20 (CI) / 3.12 (local repro)
- OS: Ubuntu 24.04.4 LTS
- XPU available: Yes
- XPU build version: 20260000
Contributor guide
Assessment
This issue has not been assessed yet.