[bugfix] skip full-size cat copy when single shard (TP==1) to avoid OOM

[ys2025-AI]

What does this PR do?

This PR fixes an OOM (out-of-memory) issue that occurs during colocate sync when tensor parallelism size is 1 (single shard).

Problem

In _get_weight, torch.concat(tensor, dim=0) was always called regardless of the number of shards. When there is only a single shard (TP==1), this still allocates a new full-size tensor copy, causing unnecessary peak GPU memory spikes that lead to OOM in the colocate sync path.

Fix

Skip the torch.concat copy when the shard list contains a single element:

• For tensor: directly use tensor[0] when len(tensor) == 1, otherwise concatenate as before.
• For mg_scale_inv: apply the same single-shard shortcut (only when it is a list/tuple of length 1).

This removes one redundant full-size tensor copy in the TP==1 path, reducing peak GPU memory usage.

Changes

• src/mcore_bridge/bridge/gpt_bridge.py: updated the shard-merging logic in _get_weight to short-circuit on single-shard inputs.

[gemini-code-assist]

Code Review

This pull request optimizes tensor concatenation in gpt_bridge.py by skipping torch.concat when there is only a single shard (length of 1) to prevent Out-Of-Memory (OOM) errors during colocate sync. The reviewer suggested simplifying the handling of mg_scale_inv by removing a redundant type check, as it is already guaranteed to be a list if not None.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.