Create option to control data type for tensor parallel all-reduce

[veritas9872]

Is your feature request related to a problem? Please describe.

I have found that using BF16 for Tensor Parallel all-reduce causes major differences in the logit values when using different TP values. This is because using BF16 buffers for TP transfers, while reducing the memory, causes errors to accumulate with each layer.

However, this can cause lots of problems when logit values are important. Currently, an average error of 15% is not uncommon when using different TP levels with BF16 buffers. This error is negligible when using FP32 buffers.

Also, now that async-TP has been introduced, I think that using FP32 for TP data transfer will not have the same effect on latency as before.

Describe the solution you'd like

Many users may prefer using BF16 for TP data transfer to reduce latency. However, having an option for FP32 TP/CP model parallel would be great for accuracy sensitive users. Also, this would allow easier changing of TP degree during training, etc.

Describe alternatives you've considered

N/A.

Additional context

From the code in the TE PyTorch Linear module, I am not 100% certain if BF16 is being used for TP/CP data transfer. However, I am reasonably certain that this is the case since no documentation has been provided on the issue.

[veritas9872]

Output logits for Llama 3.2 1B Instruct using NeMo 2.

Results for TP=1

HuggingFace outputs at the top, NeMo 2 outputs at the bottom.

Results for TP=8

NeMo 2 outputs with TP=8.

import os

import torch
import torch.distributed as dist
from nemo.collections import llm
from nemo import lightning as nl
from nemo.lightning import io
from nemo.lightning.ckpt_utils import ckpt_to_context_subdir
from transformers import AutoModelForCausalLM, AutoTokenizer


def convert():
    model_name = "meta-llama/Llama-3.2-1B-Instruct"
    nemo_model_path = "nemo_llama_32_1b_instruct"

    llm.import_ckpt(
        model=llm.LlamaModel(llm.Llama32Config1B()),
        source=f"hf://{model_name}",
        output_path=nemo_model_path,
    )


@torch.inference_mode()
def check(text: str):
    model_name = "meta-llama/Llama-3.2-1B-Instruct"
    nemo_model_path = "nemo_llama_32_1b_instruct"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_ids = tokenizer(
        text=text,
        add_special_tokens=False,
        return_attention_mask=False,
        return_tensors="pt",
    ).input_ids

    device = torch.device("cuda:2")
    with torch.cuda.device(device):
        hf_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float32,
        )
        hf_model.tie_weights()
    hf_model.config.use_cache = False
    hf_model.eval()
    hf_model.to(device)

    strategy = nl.MegatronStrategy(
        tensor_model_parallel_size=1,
    )
    trainer = nl.Trainer(
        devices=[1],
        accelerator="gpu",
        strategy=strategy,
        plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"),
    )
    trainer.strategy._setup_optimizers = False
    fabric = trainer.to_fabric()
    model: io.TrainerContext = io.load_context(
        path=ckpt_to_context_subdir(nemo_model_path),
        subpath="model",
    )
    model = fabric.load_model(nemo_model_path, model).cuda(1)
    model = model.module
    model.eval()

    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        hf_output = hf_model(input_ids=input_ids.to(device)).logits

    position_ids = torch.arange(input_ids.size(1), dtype=torch.int64, device=1)
    nemo_output = model(
        input_ids=input_ids.cuda(1),
        attention_mask=None,
        position_ids=position_ids,
    )

    breakpoint()


@torch.inference_mode()
def check_tp(text: str):
    dist.init_device_mesh(
        device_type="cuda",
        mesh_shape=(8,),
        mesh_dim_names=("tp",),
    )
    local_rank = int(os.getenv("LOCAL_RANK", 0))
    model_name = "meta-llama/Llama-3.2-1B-Instruct"
    nemo_model_path = "nemo_llama_32_1b_instruct"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_ids = tokenizer(
        text=text,
        add_special_tokens=False,
        return_attention_mask=False,
        return_tensors="pt",
    ).input_ids

    strategy = nl.MegatronStrategy(
        tensor_model_parallel_size=8,
    )
    trainer = nl.Trainer(
        devices=8,
        accelerator="gpu",
        strategy=strategy,
        plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"),
    )
    trainer.strategy._setup_optimizers = False
    fabric = trainer.to_fabric()
    model: io.TrainerContext = io.load_context(
        path=ckpt_to_context_subdir(nemo_model_path),
        subpath="model",
    )
    model = fabric.load_model(nemo_model_path, model).cuda(local_rank)
    model = model.module
    model.eval()

    position_ids = torch.arange(
        input_ids.size(1),
        dtype=torch.int64,
        device=local_rank,
    )
    nemo_output = model(
        input_ids=input_ids.cuda(local_rank),
        attention_mask=None,
        position_ids=position_ids,
    )

    if local_rank == 0:
        breakpoint()
    else:
        dist.barrier()

[veritas9872]

Commands:

python -m fire debug.py convert

python -m fire debug.py check --text "Summer is warm. Winter is cold."

torchrun --nproc-per-node=8 --nnodes=1 --standalone -m fire debug.py check_tp --text "Summer is warm. Winter is cold."

[veritas9872]

I used the NeMo 25.04 NGC container for my experiments.
The prompt "Summer is warm. Winter is cold." was chosen because it is known to cause massive activations (https://openreview.net/forum?id=F7aAhfitX6#discussion).