Skip to content

Fix device migration for Sionna PHY module state#1168

Open
DogWY wants to merge 1 commit into
NVlabs:mainfrom
DogWY:fix/module-device-migration-state
Open

Fix device migration for Sionna PHY module state#1168
DogWY wants to merge 1 commit into
NVlabs:mainfrom
DogWY:fix/module-device-migration-state

Conversation

@DogWY

@DogWY DogWY commented Jun 12, 2026

Copy link
Copy Markdown

Fixes #1167

Description

This PR fixes device migration behavior for sionna.phy modules that keep logical device state, cached tensors, or precomputed tensor state outside PyTorch's normal parameter/buffer migration path.

Previously, after constructing a PHY module on CPU and then calling .to(cuda_device), some internal state could remain on the original device. This could leave Object.device stale, keep cached tensors on the wrong device, or produce CPU outputs for CUDA inputs in affected modules.

The fix makes device-dependent internal state follow PyTorch module migration semantics while preserving trainable torch.nn.Parameter inputs where applicable.

Changes

  • Derive Object.device from a non-persistent _device_ref buffer instead of a plain string field.
  • Register precomputed and cached tensors as non-persistent buffers across affected PHY components.
  • Convert Trellis into a torch.nn.Module so its transition tensors participate in .to(...) migration.
  • Preserve trainable Parameter inputs for custom constellation points, filter coefficients, and window coefficients.
  • Update affected areas including channel, mapping, signal, MIMO, OFDM, FEC, and NR components.
  • Add regression tests for device migration and trainable-parameter preservation.

Tests

  • git diff --check main..HEAD
  • Full unit test suite was previously run on an earlier version of this fix on a CUDA-enabled Ubuntu machine:
    • OS: Ubuntu 24.04.3 LTS
    • GPU: NVIDIA RTX A4000
    • NVIDIA driver: 535.288.01
    • Python: 3.12.13
    • PyTorch: 2.11.0+cu126
    • pytest: 9.0.3
    • Result: 3901 passed, 34 skipped, 40 warnings in 6:32:37

Signed-off-by: DogWY <wangyan04221@outlook.com>
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.

sionna.phy modules do not fully migrate logical device and tensor state with PyTorch .to(device)

1 participant