quanto + musubi block swap compatibility

[bghira]

This pull request enhances device management for quantized tensors in the Musubi block swap helper and introduces comprehensive tests for these scenarios. The main improvements are the addition of specialized logic to correctly detect, move, and verify "quanto" quantized tensors within modules, and the introduction of a test suite to ensure correct behavior, especially when standard PyTorch methods cannot be used.

Device management for quantized tensors:

• Added helper functions (_same_device, _is_quanto_tensor, _tensor_on_device, _module_has_quanto_tensor, _move_quanto_tensor_to_device, _move_module_without_swapping_quanto_params) to detect and recursively move "quanto" quantized tensors and their sub-tensors to the correct device, ensuring all relevant attributes are properly transferred.
• Modified the _move_module method to use the custom quanto-aware logic when quantized tensors are present, falling back to the default .to(device) otherwise.

Testing improvements:

• Added a new test suite in tests/test_musubi_block_swap.py to verify correct streaming of quantized modules, including handling of QLinear layers and fallback logic when the standard _apply method fails.

These changes ensure robust support for quantized models and improve test coverage for device transfer logic.

[Copilot]

Pull request overview

This pull request updates the Musubi block swap helper to correctly detect and move Optimum Quanto quantized tensors (including their internal backing tensors) across devices without relying on Module.to()/_apply, and adds a targeted unit test to validate the behavior with QLinear when _apply is unusable.

Changes:

• Added Quanto-aware device comparison and tensor/module transfer helpers to safely move quantized weights and their internal tensors.
• Updated MusubiBlockSwapManager._move_module() to use the custom transfer path when Quanto quantized tensors are present, otherwise preserving the existing .to(device) behavior.
• Added a new unittest validating block swap streaming for a Quanto-quantized module with QLinear._apply forced to fail.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
simpletuner/helpers/musubi_block_swap.py	Adds Quanto tensor detection + custom, recursion-safe device transfer and uses it during block streaming when needed.
tests/test_musubi_block_swap.py	Introduces a regression test ensuring quantized blocks can stream in/out even when QLinear._apply errors.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.