[BUG] Accuracy fluctuation with tensor parallel on different card number

Describe the bug
issue 1. tensor_parallel=2 returns differnt result comparing the output without tensor parallel. It exists since deepspeed==0.16.0.
issue 2. deepspeed==0.17.4 returns a new output for tensor_parallel=2 compared to deepspeed==0.16.0 and still have issue 1.

To Reproduce
Steps to reproduce the behavior:

1. start a docker on H20, Docker: nvcr.io/nvidia/pytorch:25.02-py3
2. run below script with tp=2 and no tp. cmd: deepspeed --include="localhost:2,3" test_ds.py
3. compare the results

# test_ds.py
import transformers
import deepspeed
import os

import random
import numpy as np
import torch

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

try:
    world_size = torch.distributed.get_world_size()
    local_rank = torch.distributed.get_rank()
except:
    local_rank = int(os.getenv("LOCAL_RANK", "-1"))
    world_size = int(os.getenv("WORLD_SIZE", "-1"))


# model_name_or_path = '/ssd/xinhe/Llama-3.1-8B-Instruct/'
model_name_or_path = 'Llama-3.1-8B-Instruct'
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path)
config = transformers.AutoConfig.from_pretrained(model_name_or_path)
# using memory mapping with torch_dtype=config.torch_dtype
model = transformers.AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=config.torch_dtype)
from neural_compressor.torch.utils import local_rank, logger, world_size
if world_size > 1:
    ds_inference_kwargs = {
        "dtype": config.torch_dtype,
        "tensor_parallel": {"tp_size": world_size},
        #"keep_module_on_host": True,
    }

    ds_model = deepspeed.init_inference(model, **ds_inference_kwargs)
    model = ds_model.module

model = model.eval().to("cuda").to(torch.bfloat16)
with torch.no_grad():
    out = model(torch.tensor([[10,20,30]]).to('cuda'))[0]
print(out)
'''
** deepspeed==0.16.0 (w/o lm_head all-reduce) **
1 card (output_case 1):
tensor([[[-18.5000, -13.1250, -13.0625,  ...,   6.7500,   6.7500,   6.7500],
         [  4.2500,   5.9688,   4.7500,  ...,  -6.4062,  -6.4062,  -6.4062],
         [ -1.2656,  -7.8125,   3.3438,  ...,  -3.0000,  -3.0000,  -3.0000]]],
       device='cuda:0', dtype=torch.bfloat16)
2 card (output_case 2, different with case 1):
tensor([[[-18.8750, -13.3125, -13.2500,  ...,   6.8750,   6.8750,   6.8750],
         [  4.2188,   5.9688,   4.7188,  ...,  -6.4062,  -6.4062,  -6.4062],
         [ -1.2422,  -7.7812,   3.3438,  ...,  -2.9688,  -2.9688,  -2.9688]]],
       device='cuda:0', dtype=torch.bfloat16)

deepspeed==0.17.4 (w/ lm_head all-reduce)
1 card (same with output_case 1):
tensor([[[-18.5000, -13.1250, -13.0625,  ...,   6.7500,   6.7500,   6.7500],
         [  4.2500,   5.9688,   4.7500,  ...,  -6.4062,  -6.4062,  -6.4062],
         [ -1.2656,  -7.8125,   3.3438,  ...,  -3.0000,  -3.0000,  -3.0000]]],
       device='cuda:0', dtype=torch.bfloat16)
2 card (output_case 3, different with case 1 and case 2):
tensor([[[-18.7500, -13.3125, -13.2500,  ...,   6.8750,   6.8750,   6.8750],
         [  4.2500,   5.9375,   4.7188,  ...,  -6.4062,  -6.4062,  -6.4062],
         [ -1.2422,  -7.8125,   3.3594,  ...,  -2.9688,  -2.9688,  -2.9688]]],
       device='cuda:0', dtype=torch.bfloat16)
'''

Expected behavior
It should provide a consistent result or let me know why it happens if it cannot be resolved.

ds_report output
Please run ds_report to give us details about your setup.

System info (please complete the following information):

• Docker: nvcr.io/nvidia/pytorch:25.02-py3
• H20 GPU

Docker context
Docker: nvcr.io/nvidia/pytorch:25.02-py3

The reproduction is in test_ds.py and uses deepspeed.init_inference with tensor_parallel; start by running the stated deepspeed command with one and two cards. Compare outputs under DeepSpeed 0.16.0 and 0.17.4, and review the supplied ds_report output. Done means the discrepancy is explained or consistent behavior is established.