Skip to content

Conversation

@SongXiaoXi
Copy link
Collaborator

This pull request introduces a minor import adjustment and a fix to device handling when pinning memory for CUDA tensors. The main changes ensure that memory pinning happens on the correct CUDA device and resolve an import path issue.

CUDA device handling improvements:

  • In checkpoint_engine/pin_memory.py, the current CUDA device is now explicitly retrieved and set (device_index = torch.cuda.current_device() and torch.cuda.set_device(device_index)) before pinning memory, ensuring that pinning occurs on the correct device. [1] [2]

Import path correction:

  • In examples/update.py, the import for request_inference_to_update is fixed by importing it directly from checkpoint_engine instead of checkpoint_engine.ps, resolving a potential import error.

@blahgeek blahgeek merged commit 4a73109 into MoonshotAI:main Jan 9, 2026
2 checks passed
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.

2 participants