Skip to content

Conversation

@bghira
Copy link
Owner

@bghira bghira commented Jan 29, 2026

This pull request makes targeted improvements to how text encoders and linear layers are managed when using RamTorch, particularly around device placement and partial replacement scenarios. The changes ensure better compatibility and correctness when RamTorch is enabled, especially in mixed-device and percentage-based replacement cases.

Device management and RamTorch integration:

  • Improved logic in move_text_encoders (in factory.py) to always skip moving text encoders when RamTorch is enabled, regardless of the target device, to support both partial and full RamTorch setups.

Partial RamTorch replacement handling:

  • In replace_linear_layers_with_ramtorch (in ramtorch.py), added logic to track remaining eligible layers when using percentage-based replacement, ensuring only the correct subset is replaced.
  • Enhanced handling so that, after percentage-based replacement, any remaining eligible nn.Linear layers (or those matching specific patterns) are moved to the resolved device if they are still on CPU, maintaining device consistency.

@bghira bghira merged commit 5638c97 into main Jan 29, 2026
2 checks passed
@bghira bghira deleted the bugfix/ramtorch-percentage-text-enc branch January 29, 2026 18:04
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