[wip] split to_device into collection + execution

[jameszianxuTT]

Ticket

One possible solution to #3218. While a local patch could resolve this for now, I think the investigation in #3218 shows that we need to be more careful when moving aliased objects to device.

Problem description

Provide context for the problem.

What's changed

Only move objects .to(device) once. This prevents an object and its aliases from being moved .to(device) separately creating separate objects, leading to broken shard spec application

Checklist

• New/Existing tests provide coverage for changes

[jameszianxuTT]

Would like to think a bit more on whether this is a good solution, since it adds a chunk of complexity to the nice .to_device() that we had before. However, these kinds of bugs are subtle, so it might be worth catching

Another alternative is having this alias set up in a different way in tests/torch/models/kimi_k2/utils.py

    @property
    def keys(self):
        """Return compressed_kv as keys. This ensures keys is always the same object as compressed_kv."""
        return self.compressed_kv

    @keys.setter
    def keys(self, value):
        """Set compressed_kv when keys is assigned. This keeps them in sync."""
        self.compressed_kv = value

    @property
    def values(self):
        """Return k_pe as values. This ensures values is always the same object as k_pe."""
        return self.k_pe

    @values.setter
    def values(self, value):
        """Set k_pe when values is assigned. This keeps them in sync."""
        self.k_pe = value

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 28.33%. Comparing base (dfa3281) to head (e0ffbc9).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3274   +/-   ##
=======================================
  Coverage   28.33%   28.33%           
=======================================
  Files          33       33           
  Lines        4094     4094           
=======================================
  Hits         1160     1160           
  Misses       2934     2934           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.