The WeightSynchronizer FFI derived the host staging slot as shard_idx % num_shards, which only works if a slice's device ordinals are contiguous and aligned to a multiple of num_shards. Device ordinals are not guaranteed to be, so pass the staging slot explicitly as a range-checked logical_shard_idx input instead of deriving it from the device ordinal.#385
Open
copybara-service[bot] wants to merge 1 commit into
Open
The WeightSynchronizer FFI derived the host staging slot as shard_idx % num_shards, which only works if a slice's device ordinals are contiguous and aligned to a multiple of num_shards. Device ordinals are not guaranteed to be, so pass the staging slot explicitly as a range-checked logical_shard_idx input instead of deriving it from the device ordinal.#385copybara-service[bot] wants to merge 1 commit into
shard_idx % num_shards, which only works if a slice's device ordinals are contiguous and aligned to a multiple of num_shards. Device ordinals are not guaranteed to be, so pass the staging slot explicitly as a range-checked logical_shard_idx input instead of deriving it from the device ordinal.#385copybara-service[bot] wants to merge 1 commit into
