[bugfix] fix ZCH finetune from checkpoint with different world size

[tiankongdeguiji]

Summary

Loading a ZCH (managed-collision) checkpoint into a model with a different world size crashed in MCHManagedCollisionModule.validate_state():

AssertionError: shard within range [500000, 1000000] cannot be built out of
segements tensor([0, 1000000, -1, ..., -1, -1, -1])

Two distinct issues stacked on top of each other:

1. _output_segments_tensor is a fixed-shape [1025] replicated buffer in MCHManagedCollisionModule whose contents are rank-specific partition boundaries. Its shape is identical across world sizes, so torch DCP silently overwrites the freshly-built local value with the saved one, then validate_state() fails because _output_global_offset (a Python int that is not in the state dict) no longer appears in the segments.

2. The per-position MCH lookup state (_mch_sorted_raw_ids, _mch_remapped_ids_mapping, _mch_<metadata>) is wrapped as ShardedTensor by ShardedManagedCollisionCollection._initialize_torch_state (torchrec mc_modules.py:262). Across matching world sizes this works fine, but across different world sizes torch DCP only does byte-level position slicing — and the values in _mch_remapped_ids_mapping are global slot indices in [output_global_offset, output_global_offset + zch_size), so the loaded values fall outside the new local range and crash the FBGEMM TBE bounds check / CUDA gather kernel during the first training step.

Fix

• PartialLoadPlanner always skips _output_segments_tensor; the local buffer is rebuilt by fix_mch_state (extended to handle the post-init_parameters all-zeros case used by the export path) and is also called from restore_model so the load post-hook validation passes.
• When the saved sharding plan's world size differs from the current world size, PartialLoadPlanner also skips the sharded MCH state buffers, and restore_model runs a new _redistribute_mch_state pass that:
  1. Reads the full saved global tensors for each MC table.
  2. For every non-empty saved entry, computes target_rank = saved_remapped_value // local_zch_size. This matches the row-wise position-based sharding that ShardedTensor already uses for the embedding weight, so the (raw_id → embedding row) binding is preserved end-to-end across the world size change.
  3. Writes the kept entries into the local module's buffers, fills the unused positions with the local-range remapped values, and calls _sort_mch_buffers to keep the binary-search invariant.

The embedding weight tensors are still loaded by torch DCP via the normal ShardedTensor path; the value-based bucket above is chosen to align with that position-based slice, so no extra weight movement is needed.

Test plan

• New test_multi_tower_din_zch_finetune_world_size_change integration test in tzrec/tests/rank_integration_test.py: trains the ZCH multi-tower DIN config on 1 GPU, then runs --fine_tune_checkpoint on 2 GPUs.
• Existing test_multi_tower_din_zch_with_fg_train_eval_export (same world size, ZCH train + export) still passes — exercises the fix_mch_state export path.
• Existing test_multi_tower_din_fg_encoded_finetune (non-ZCH finetune) still passes — sanity check that the planner skip does not affect non-MCH modules.
• Manual end-to-end: 1-GPU → 2-GPU finetune and 2-GPU → 1-GPU finetune both reach Train and Evaluate Finished. cleanly.

🤖 Generated with Claude Code

[github-actions]

Bug risk: broad except Exception silently disables MCH redistribution

If the plan file exists but is corrupted or has an unexpected schema, this catch-all swallows the error and leaves saved_world_size = None, which sets needs_mch_redistribution = False. The code then falls back to position-based ShardedTensor loading — exactly the broken path this PR is fixing — without any indication that redistribution was needed but couldn't be determined.

Consider distinguishing "file not found" (acceptable) from "file exists but unparseable" (should raise or at least log at ERROR level):

Suggested change

	except Exception as e:
	logger.warning(f"Failed to inspect saved plan {plan_path}: {e}")
	except (json.JSONDecodeError, KeyError, TypeError) as e:
	logger.error(
	f"Plan file {plan_path} exists but failed to parse: {e}. "
	"MCH redistribution will be skipped — this may cause "
	"incorrect state loading if the world size has changed."
	)

[github-actions]

Robustness: add a defensive assertion for duplicate remapped IDs

If a corrupted checkpoint contains duplicate values in _mch_remapped_ids_mapping, taken would have fewer True entries than expected, making unused too short. new_remapped[n:] would then be partially filled with stale arange values — silently corrupting the model state without any error.

Consider adding:

assert kept_values.unique().numel() == kept_values.numel(), (
    f"MCH [{prefix}]: duplicate remapped IDs in checkpoint data"
)

And after line 492:

assert unused.numel() >= local_zch_size - n, (
    f"MCH [{prefix}]: insufficient unused slots ({unused.numel()}) "
    f"to fill {local_zch_size - n} empty positions"
)

[github-actions]

Nit: hardcoded magic number 1025

This value presumably matches a torchrec internal constant (MAX_WORLD_SIZE + 1 or similar). Consider deriving it from the existing buffer's shape to be resilient against future torchrec changes:

Suggested change

	output_segments + [-1] * (1025 - len(output_segments)),
	output_segments + [-1] * (buf.shape[0] - len(output_segments)),

This also handles the is_meta case correctly since meta tensors preserve their shape.

[github-actions]

Nit: ambiguous zch_size in docstring

Here zch_size means the global total across all ranks, but throughout the code m._zch_size (assigned to local_zch_size) is the per-rank size. Consider clarifying:

Suggested change

	routes inputs to the rank that owns each value's range. Position-based
	`ShardedTensor` resharding therefore does not preserve semantics across
	slot indices in `[0, global_zch_size)` (where `global_zch_size = local_zch_size *
	world_size`) — NOT positions, NOT bytes — and torchrec

[github-actions]

Code Review Summary

Well-crafted bugfix for a genuinely tricky distributed systems problem. The value-aware MCH redistribution algorithm is sound, the inline comments are excellent (especially explaining why position-based ShardedTensor slicing breaks across world sizes), and the second commit narrowing the condition to cur > saved shows good iterative debugging. A few items worth addressing:

Issues

1. Silent fallback on plan file corruption — The broad except Exception when parsing the saved plan file silently disables MCH redistribution if the file is corrupt, falling back to the exact broken path this PR fixes. See inline comment.

2. No defense against duplicate remapped IDs from corrupted checkpoints — taken[kept_values - local_offset] = True with duplicate values would produce fewer unused slots than expected, silently corrupting the local buffer. See inline comment.

3. Hardcoded magic number 1025 in fix_mch_state — should derive from the buffer's existing shape (buf.shape[0]). See inline comment.

Testing gaps

4. Integration test only covers 1→2 GPU direction. The reverse (2→1) exercises a different code path (fix_mch_state rebuild without redistribution) that has no dedicated test. The PR description mentions manual verification of both directions — consider adding the 2→1 test as well.

5. Test verifies non-crash but not correctness. The assertions check self.success and steps > 0, but don't verify that embeddings were actually preserved across the world-size change. A smoke check comparing a known feature's embedding before and after would strengthen confidence in the redistribution logic.

Scaling note

6. Every rank independently loads the full global MCH state in _read_full_mch_tensors. For a bugfix this is fine, but at larger scale (many ranks, large ZCH tables) this multiplies both I/O and CPU memory by world_size. Worth a # TODO for future optimization (rank-0 load + broadcast).

Positive notes

• The algorithm in _redistribute_mch_state correctly aligns value-based MCH bucketing with the position-based ShardedTensor slicing used for embedding weights — this is the key insight that makes the fix work end-to-end.
• The PartialLoadPlanner skip logic is clean and well-documented.
• The second commit narrowing needs_mch_redistribution to only cur > saved is a good catch that preserves the ckpt_param_map_path remap path for export.

🤖 Generated with Claude Code

[github-actions]

Bug (deadlock risk): If _ckpt_world_size raises on some ranks but not others (e.g. NFS latency, partial file visibility), those ranks return False here while others return True. Since _redistribute_mch_state calls dist.all_to_all_single — a collective — the process group will deadlock if only a subset of ranks enters.

Consider synchronizing the decision across ranks after computing it locally:

needs = _needs_mch_redistribution(model_ckpt_path, cur_world_size)
if dist.is_initialized():
    flag = torch.tensor([int(needs)], device="cuda")
    dist.all_reduce(flag, op=dist.ReduceOp.MAX)
    needs = flag.item() > 0

Or alternatively, let the exception propagate (remove the try/except) so all ranks fail consistently rather than silently diverging.

[github-actions]

Bug (silent data corruption): remapped_recv - local_offset is used as an index into taken (size cps) with no bounds check. If any received value falls outside [local_offset, local_offset + cps), PyTorch's negative-index wrapping will silently corrupt unrelated positions rather than raising an error.

Add a bounds check before indexing:

indices = remapped_recv - local_offset
if n > 0:
    assert (indices >= 0).all() and (indices < cps).all(), (
        f"MCH [{prefix}] rank {cur_rank}: received remapped values outside "
        f"local range [{local_offset}, {local_offset + cps})"
    )
    taken[indices] = True

[github-actions]

Nit: assert is stripped under python -O. For a runtime invariant guarding data integrity in distributed code, prefer if ... : raise RuntimeError(...).

Suggested change

	assert n <= cps, (
	f"MCH [{prefix}] rank {cur_rank}: received {n} > local zch_size {cps}"
	)
	if n > cps:
	raise RuntimeError(
	f"MCH [{prefix}] rank {cur_rank}: received {n} > local zch_size {cps}"
	)

[github-actions]

Suggestion: max(ranks) + 1 assumes contiguous rank numbering. If a checkpoint is partially written (e.g. rank 1 shard missing, only ranks 0 and 2 present), this silently returns the wrong world size, which could cascade into incorrect redistribution decisions.

Consider adding a consistency check:

expected = max(ranks) + 1
if len(ranks) != expected:
    raise RuntimeError(
        f"Checkpoint {ckpt_dir} has rank gaps: found {sorted(ranks)}, "
        f"expected contiguous 0..{expected - 1}"
    )

[github-actions]

Fragility warning: This sequence relies on model.state_dict() (line 487) returning tensors that share underlying storage with model._buffers, so that _redistribute_mch_state's .copy_() into the buffers also updates what state_dict references. If PyTorch/TorchRec ever changes state_dict() to return detached copies, the redistributed values would be silently overwritten by pre-redistribution data in load_state_dict.

Consider explicitly copying the redistributed buffer values back into state_dict after redistribution, or adding a comment + assertion that the storage sharing invariant holds (e.g. assert state_dict[some_key].data_ptr() == model.some_buffer.data_ptr()).

[github-actions]

Code Review Summary

Well-crafted fix for a gnarly problem — the value-aware all_to_all_single redistribution is the right approach, and the commit history shows careful iterative refinement of the divisibility condition and skip scope. The PR description and commit messages are exemplary.

Issues to address

1. Deadlock risk from inconsistent redistribution decision across ranks (comment)
_needs_mch_redistribution catches filesystem exceptions and returns False. If _ckpt_world_size fails on some ranks but not others (NFS latency, partial file visibility), some ranks enter all_to_all_single while others skip it → distributed deadlock. Either synchronize the decision via all_reduce or let the exception propagate.

2. Missing bounds check on remapped_recv indexing (comment)
taken[remapped_recv - local_offset] = True has no guard. Out-of-range values silently corrupt via PyTorch negative-index wrapping rather than raising. Add an assertion before indexing.

3. assert used for runtime invariant (comment)
The assert n <= cps check is stripped under python -O. Use if/raise RuntimeError for data integrity checks in distributed code.

4. _ckpt_world_size assumes contiguous rank numbering (comment)
max(ranks) + 1 silently returns the wrong world size if any intermediate shard file is missing. A len(ranks) != max(ranks) + 1 check would catch corrupted checkpoints early.

5. Implicit storage-sharing dependency in restore_model (comment)
The state_dict → load → redistribute → load_state_dict flow relies on state_dict tensors sharing storage with model buffers. If this invariant breaks in a future PyTorch/TorchRec version, redistributed values would be silently overwritten.

Optimization opportunity

Fuse multiple all_to_all_single calls per MC module. Currently K+2 collectives are issued per module (raw_ids + remapped + K metadata buffers). Since all share identical split sizes, they could be concatenated into one tensor, exchanged with a single all_to_all_single, and sliced apart. This reduces collective synchronization overhead, especially at higher rank counts.

Testing

The integration test covers the critical 1→2 GPU path and is correctly gated on GPU availability. Consider adding unit tests for the pure helper functions (_ckpt_world_size, _needs_mch_redistribution, _strip_dmp_prefix) — they require no GPU/distributed setup and would catch edge cases like non-contiguous ranks or the divisibility boundary conditions.

🤖 Generated with Claude Code