Skip to content

fix(numeric): support legacy PyTorch autocast fallback - #185

Open
PinkTulips139 wants to merge 2 commits into
Tencent:mainfrom
PinkTulips139:fix-legacy-torch-autocast-compat
Open

fix(numeric): support legacy PyTorch autocast fallback#185
PinkTulips139 wants to merge 2 commits into
Tencent:mainfrom
PinkTulips139:fix-legacy-torch-autocast-compat

Conversation

@PinkTulips139

Copy link
Copy Markdown
Contributor

Summary

Adds a backward-compatible autocast-disabled context for legacy PyTorch versions that do not provide the top-level torch.autocast API.

The issue was observed in the repository's Python 3.8 / PyTorch 1.8 CI environment, where existing Mixture, MoA, and MoT tests failed with AttributeError: module 'torch' has no attribute 'autocast'.

Changes

  • Preserve the existing torch.autocast(device_type=..., enabled=False) path on modern PyTorch.
  • Fall back to torch.cuda.amp.autocast(enabled=False) for legacy CUDA environments.
  • Use nullcontext() for legacy CPU, MPS, and unsupported device types.
  • Keep compatibility handling centralized in disabled_autocast() rather than patching individual Mixture modules.
  • Add regression tests for modern CPU/CUDA/MPS paths, legacy CUDA, legacy CPU/MPS, and unknown devices.

Scope

This PR does not change:

  • routing behavior;
  • model outputs;
  • optimizer logic;
  • training configuration;
  • LoRA behavior;
  • CI workflow definitions.

It only provides a compatibility fallback for the shared numerical helper.

Validation

  • 8 passed in tests/test_torch_legacy_compat.py
  • 13 passed in the previously affected Mixture test files
  • Ruff passed
  • Python compilation passed
  • git diff --check passed

No GPU is required for the regression tests; the legacy CUDA path is validated with mocks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant