intel / intel/torch-xpu-ops

[release/2.14] Muon optimizer state_dict cross-device numeric tolerance failure

Open
#4,263 1 comment 0 reactions 1 assignee Claimed by @AKloniecki View on GitHub
os: Linux os: Windows skipped test: ut
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.