第 12 章：多机多卡调试地狱

export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL
export TORCH_DISTRIBUTED_DEBUG=DETAIL

# 检查节点间网络连通性
ping <other_node_ip>
nc -zv <other_node_ip> 29500  # 检查端口是否开放

# 检查 GPU 可见性
nvidia-smi -L
echo $CUDA_VISIBLE_DEVICES

# 单机测试
python -m torch.distributed.run --nproc_per_node=8 --nnodes=1 test_nccl.py

# 多机测试（在主节点运行）
python -m torch.distributed.run \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr=<master_ip> \
    --master_port=29500 \
    test_nccl.py

RuntimeError: NCCL error in: ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1191
unhandled system error, NCCL version 2.14.3

# 1. 确保 GPU 编号正确
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

# 2. 增加共享内存
docker run --shm-size=32gb ...  # Docker 环境
# 或修改系统配置
echo "kernel.shmmax = 68719476736" >> /etc/sysctl.conf
echo "kernel.shmall = 4294967296" >> /etc/sysctl.conf
sysctl -p

# 3. 降级或升级 NCCL
pip install torch==2.0.1+cu118  # 使用兼容版本

torch.distributed.DistBackendError: NCCL timeout after 1800 seconds

# 1. 增加超时时间
export NCCL_TIMEOUT=3600  # 单位：秒
export NCCL_ASYNC_ERROR_HANDLING=1

# 2. 使用更快的网络协议
export NCCL_IB_DISABLE=0  # 启用 InfiniBand
export NCCL_SOCKET_IFNAME=eth0  # 指定网络接口

# 3. 检查进程状态
ps aux | grep python  # 查看是否有僵尸进程
htop  # 查看 CPU/内存使用情况

NCCL WARN Unhandled System Error while waiting for event

# 1. 检查 PCIe 和 NVLink 状态
nvidia-smi topo -m

# 2. 优化 P2P 通信
export NCCL_P2P_DISABLE=0
export NCCL_P2P_LEVEL=NVL  # 使用 NVLink

# 3. 设置合理的 NCCL 树形拓扑
export NCCL_TREE_THRESHOLD=0  # 总是使用树形算法

# 检查 IB 状态
ibstat
ibping -S  # 在一个节点启动服务器
ibping -c <server_guid>  # 在另一个节点测试

# NCCL IB 配置
export NCCL_IB_HCA=mlx5_0,mlx5_1  # 指定 HCA
export NCCL_IB_GID_INDEX=3  # RoCE 环境需要
export NCCL_IB_TC=106  # Traffic Class
export NCCL_IB_SL=0  # Service Level

# 1. 选择正确的网络接口
export NCCL_SOCKET_IFNAME=eth0,eth1  # 可以指定多个
export NCCL_SOCKET_NTHREADS=8  # 增加 socket 线程数
export NCCL_NSOCKS_PERTHREAD=4  # 每线程 socket 数

# 2. TCP 缓冲区优化
echo "net.core.rmem_max = 134217728" >> /etc/sysctl.conf
echo "net.core.wmem_max = 134217728" >> /etc/sysctl.conf
echo "net.ipv4.tcp_rmem = 4096 87380 134217728" >> /etc/sysctl.conf
echo "net.ipv4.tcp_wmem = 4096 65536 134217728" >> /etc/sysctl.conf
sysctl -p

# 调试相关
NCCL_DEBUG=INFO/WARN/ERROR  # 日志级别
NCCL_DEBUG_FILE=/path/to/log  # 日志输出文件

# 性能调优
NCCL_BUFFSIZE=8388608  # 缓冲区大小（默认 4MB）
NCCL_NTHREADS=512  # NCCL 线程数
NCCL_MAX_NCHANNELS=16  # 最大通道数

# 网络选择
NCCL_NET_GDR_LEVEL=5  # GPUDirect RDMA 级别
NCCL_CROSS_NIC=1  # 允许跨 NIC 通信

# 算法选择
NCCL_ALGO=Ring/Tree/CollNet  # 指定算法
NCCL_PROTO=LL/LL128/Simple  # 协议选择

# 1. 调整 GPU 映射，避免使用问题 GPU 对
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5  # 跳过 6,7

# 2. 降低功率上限防止过热
nvidia-smi -pl 300  # 设置功率上限 300W

# 3. 优化通信模式
export NCCL_P2P_LEVEL=PHB  # 只使用同一 PCIe 桥下的 P2P
export NCCL_ALGO=Tree  # 使用树形算法减少 P2P 依赖

# 4. 监控脚本
watch -n 1 'nvidia-smi --query-gpu=index,name,temperature.gpu,power.draw --format=csv'

# 1. Barrier 同步
torch.distributed.barrier()  # 所有进程必须到达

# 2. All-Reduce 操作
torch.distributed.all_reduce(tensor)  # 梯度同步

# 3. Broadcast 操作  
torch.distributed.broadcast(tensor, src=0)  # 参数广播

# 1. 优化器步进
optimizer.step()  # DDP 会自动同步梯度

# 2. 模型保存
if rank == 0:
    torch.save(model.state_dict(), path)
torch.distributed.barrier()  # 等待保存完成

# 3. 数据加载
dataloader = DataLoader(dataset, sampler=DistributedSampler(dataset))
# DistributedSampler 确保各进程数据不重复

if rank == 0 and epoch % 10 == 0:
    # 只有 rank 0 执行验证
    val_loss = validate(model, val_loader)
    torch.distributed.broadcast(val_loss, src=0)  # 死锁！其他进程未执行

if epoch % 10 == 0:
    if rank == 0:
        val_loss = validate(model, val_loader)
    else:
        val_loss = torch.zeros(1).cuda()
    torch.distributed.broadcast(val_loss, src=0)  # 所有进程都参与

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()  # 某些进程数据提前结束，未参与同步

# 方案 1：使用 drop_last
dataloader = DataLoader(dataset, drop_last=True, ...)

# 方案 2：填充数据
total_samples = len(dataset)
samples_per_rank = math.ceil(total_samples / world_size)
# 确保每个进程有相同数量的批次

try:
    output = model(input)
except Exception as e:
    print(f"Error on rank {rank}: {e}")
    continue  # 跳过这个批次，但其他进程还在等待同步！

try:
    output = model(input)
except Exception as e:
    print(f"Error on rank {rank}: {e}")
    # 通知所有进程出现异常
    error_flag = torch.tensor([1.0]).cuda()
    torch.distributed.all_reduce(error_flag)
    if error_flag.item() > 0:
        # 所有进程一起退出
        cleanup_distributed()
        sys.exit(1)

import functools
import torch.distributed as dist

def timeout_wrapper(func, timeout=300):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        import signal

        def timeout_handler(signum, frame):
            raise TimeoutError(f"{func.__name__} timeout after {timeout}s")

        signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(timeout)
        try:
            result = func(*args, **kwargs)
        finally:
            signal.alarm(0)
        return result
    return wrapper

# 使用
@timeout_wrapper
def training_step(batch):
    # 训练代码
    pass

import threading
import time

def monitor_thread(rank, interval=30):
    """监控线程，定期打印进程状态"""
    def run():
        step = 0
        while True:
            time.sleep(interval)
            print(f"[Rank {rank}] Heartbeat at step {step}, "
                  f"Memory: {torch.cuda.memory_allocated()/1e9:.2f}GB")
            step += 1

    thread = threading.Thread(target=run, daemon=True)
    thread.start()

# 在训练开始时启动
monitor_thread(rank)

# 安装 py-spy
pip install py-spy

# 分析挂起的进程
py-spy dump --pid <process_id>

# 生成火焰图
py-spy record -d 30 -o profile.svg --pid <process_id>

# 问题：rank 0 做额外的日志记录
if rank == 0:
    # 复杂的日志计算
    log_metrics(model, epoch, step)

torch.distributed.barrier()  # rank 0 太慢，其他进程超时

# 问题：所有进程同时写入
for rank in range(world_size):
    with open(f"log_{rank}.txt", "a") as f:
        f.write(metrics)  # 文件系统压力导致某些进程阻塞

# 1. 使用异步 I/O
import asyncio

async def async_log(data, filename):
    async with aiofiles.open(filename, 'a') as f:
        await f.write(data)

# 2. 错开 I/O 时机
if step % 100 == rank:  # 不同 rank 在不同步数记录
    save_checkpoint(model, f"ckpt_{rank}_{step}.pt")

# 3. 设置合理的超时
os.environ['TORCH_DISTRIBUTED_TIMEOUT'] = '3600'  # 1 小时

# 使用 gdb 附加到进程
gdb -p <pid>
(gdb) py-bt  # 查看 Python 调用栈

# 发现卡在：
File "torch/distributed/distributed_c10d.py", line 2838, in barrier
    work.wait()

# 添加调试代码
print(f"[Rank {rank}] Entering barrier at step {step}")
torch.distributed.barrier()
print(f"[Rank {rank}] Exited barrier")

# 发现 rank 43 未进入 barrier

# 原代码
if batch_idx > 0 and batch_idx % gradient_accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()

# 问题：rank 43 的数据少一个 batch，未执行最后一次 optimizer.step()

# 1. 确保所有 rank 执行相同次数的优化步骤
total_steps = len(dataloader) // gradient_accumulation_steps
for step in range(total_steps):
    for _ in range(gradient_accumulation_steps):
        batch = next(iter(dataloader), None)
        if batch is not None:
            loss = model(batch)
            loss.backward()

    optimizer.step()
    optimizer.zero_grad()

# 2. 添加同步检查点
if step % 100 == 0:
    # 同步检查，确保所有进程进度一致
    progress_tensor = torch.tensor([step], device='cuda')
    progress_list = [torch.zeros_like(progress_tensor) 
                    for _ in range(world_size)]
    torch.distributed.all_gather(progress_list, progress_tensor)

    if rank == 0:
        progress = [p.item() for p in progress_list]
        if len(set(progress)) > 1:
            print(f"Warning: Progress mismatch: {progress}")

# 诊断代码
import time
import torch.distributed as dist

def measure_gpu_speed(rank, model, dummy_batch):
    torch.cuda.synchronize()
    start = time.time()

    for _ in range(100):
        output = model(dummy_batch)
        loss = output.mean()
        loss.backward()

    torch.cuda.synchronize()
    elapsed = time.time() - start

    # 收集所有 GPU 的时间
    times = [torch.zeros(1).cuda() for _ in range(dist.get_world_size())]
    torch.distributed.all_gather(times, torch.tensor([elapsed]).cuda())

    if rank == 0:
        times = [t.item() for t in times]
        print(f"GPU speeds: {times}")
        print(f"Slowest/Fastest ratio: {max(times)/min(times):.2f}x")

class HeterogeneousDataLoader:
    def __init__(self, dataset, gpu_configs):
        """
        gpu_configs: {
            0: {'type': 'A100', 'memory': 80, 'batch_size': 8},
            1: {'type': 'V100', 'memory': 32, 'batch_size': 4},
        }
        """
        self.dataset = dataset
        self.gpu_configs = gpu_configs

    def get_batch_size(self, rank):
        # 根据 GPU 能力分配不同批大小
        return self.gpu_configs[rank]['batch_size']

    def create_loader(self, rank):
        batch_size = self.get_batch_size(rank)
        sampler = DistributedSampler(
            self.dataset,
            num_replicas=len(self.gpu_configs),
            rank=rank,
            shuffle=True
        )
        return DataLoader(
            self.dataset,
            batch_size=batch_size,
            sampler=sampler
        )

def adaptive_gradient_accumulation(rank, base_batch_size=8):
    """根据 GPU 显存动态调整梯度累积"""
    gpu_memory = torch.cuda.get_device_properties(rank).total_memory

    if gpu_memory > 80 * 1024**3:  # 80GB
        micro_batch_size = base_batch_size
        accumulation_steps = 1
    elif gpu_memory > 40 * 1024**3:  # 40GB
        micro_batch_size = base_batch_size // 2
        accumulation_steps = 2
    elif gpu_memory > 24 * 1024**3:  # 24GB
        micro_batch_size = base_batch_size // 4
        accumulation_steps = 4
    else:  # 16GB or less
        micro_batch_size = 1
        accumulation_steps = base_batch_size

    return micro_batch_size, accumulation_steps

def setup_mixed_precision(rank):
    """处理不同 GPU 的精度差异"""
    gpu_name = torch.cuda.get_device_name(rank)

    if 'A100' in gpu_name or 'H100' in gpu_name:
        # 支持 BF16
        dtype = torch.bfloat16
        use_bf16 = True
    else:
        # 只支持 FP16
        dtype = torch.float16
        use_bf16 = False

    # 统一精度设置
    all_use_bf16 = torch.tensor([use_bf16], dtype=torch.bool).cuda()
    dist.all_reduce(all_use_bf16, op=dist.ReduceOp.MIN)

    if all_use_bf16.item():
        return torch.bfloat16
    else:
        # 降级到 FP16
        if rank == 0:
            print("Warning: Falling back to FP16 due to hardware limitations")
        return torch.float16

class MixedPrecisionOptimizer:
    def __init__(self, optimizer, model, dtype):
        self.optimizer = optimizer
        self.model = model
        self.dtype = dtype
        self.grad_scaler = torch.cuda.amp.GradScaler(enabled=(dtype == torch.float16))

    def step(self):
        # 梯度转换到统一精度
        for param in self.model.parameters():
            if param.grad is not None:
                # 确保梯度精度一致
                param.grad = param.grad.to(dtype=torch.float32)

        # 同步前转换
        for param in self.model.parameters():
            if param.grad is not None:
                dist.all_reduce(param.grad)
                param.grad = param.grad / dist.get_world_size()

        # 优化器步进
        if self.dtype == torch.float16:
            self.grad_scaler.step(self.optimizer)
            self.grad_scaler.update()
        else:
            self.optimizer.step()

# 配置文件 heterogeneous_config.yaml
gpu_groups:
  high_tier:  # A100-80G
    ranks: [0, 1, 2, 3]
    batch_size: 8
    gradient_accumulation: 1
    precision: bfloat16

  low_tier:  # V100-32G  
    ranks: [4, 5, 6, 7]
    batch_size: 3
    gradient_accumulation: 3
    precision: float16

communication:
  # 分组通信策略
  allreduce_groups:

    - [0, 1, 2, 3]  # A100 内部先同步
    - [4, 5, 6, 7]  # V100 内部先同步

  # 跨组同步使用异步模式
  cross_group_async: true

training:
  # 使用 pipeline 并行缓解不均衡
  pipeline_parallel: true
  pipeline_stages:

    - ranks: [0, 1]  # A100 处理前面层
    - ranks: [2, 3]  # A100 处理中间层
    - ranks: [4, 5, 6, 7]  # V100 处理后面层（计算量较小）

class HeterogeneousTrainer:
    def __init__(self, config_path):
        self.config = yaml.load(open(config_path))
        self.rank = dist.get_rank()
        self.setup_gpu_group()

    def setup_gpu_group(self):
        # 确定当前 GPU 所属组
        for group_name, group_config in self.config['gpu_groups'].items():
            if self.rank in group_config['ranks']:
                self.group_name = group_name
                self.group_config = group_config
                break

        # 创建通信组
        for group_ranks in self.config['communication']['allreduce_groups']:
            group = dist.new_group(ranks=group_ranks)
            if self.rank in group_ranks:
                self.local_group = group

    def train_step(self, batch):
        # 根据组配置调整批大小
        micro_batch = self._split_batch(batch, self.group_config['batch_size'])

        total_loss = 0
        for step in range(self.group_config['gradient_accumulation']):
            with torch.cuda.amp.autocast(
                dtype=getattr(torch, self.group_config['precision'])
            ):
                loss = self.model(micro_batch[step])
                loss = loss / self.group_config['gradient_accumulation']

            loss.backward()
            total_loss += loss.item()

        # 分层同步策略
        self._hierarchical_allreduce()

        return total_loss

    def _hierarchical_allreduce(self):
        # 第一步：组内同步
        for param in self.model.parameters():
            if param.grad is not None:
                dist.all_reduce(param.grad, group=self.local_group)

        # 第二步：跨组同步（仅组长参与）
        if self.rank == self.group_config['ranks'][0]:
            for param in self.model.parameters():
                if param.grad is not None:
                    dist.all_reduce(param.grad)  # 全局同步

        # 第三步：组内广播
        for param in self.model.parameters():
            if param.grad is not None:
                dist.broadcast(param.grad, 
                             src=self.group_config['ranks'][0],
                             group=self.local_group)

#!/bin/bash
# monitor_heterogeneous.sh

while true; do
    echo "=== GPU Status at $(date) ==="

    # 收集所有节点的 GPU 信息
    for node in node1 node2; do
        echo "Node: $node"
        ssh $node "nvidia-smi --query-gpu=index,name,memory.used,memory.total,utilization.gpu,temperature.gpu,power.draw --format=csv"
    done

    # 检查速度差异
    echo "=== Training Speed ==="
    tail -n 8 training.log | grep "step_time"

    # 检查是否有 OOM
    if grep -q "out of memory" training.log; then
        echo "WARNING: OOM detected!"
        grep "out of memory" training.log | tail -n 5
    fi

    sleep 30
done

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    BackwardPrefetch,
    ShardingStrategy,
    CPUOffload
)
from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy,
    size_based_auto_wrap_policy
)

def setup_fsdp_for_vlm(model, vision_encoder, language_model):
    # 混合精度配置
    mp_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
        cast_forward_inputs=True
    )

    # 自动包装策略 - 关键！
    auto_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={
            # 视觉编码器层
            vision_encoder.__class__,
            # 语言模型层  
            type(language_model.layers[0]),
        }
    )

    # CPU Offload（显存不足时启用）
    cpu_offload = CPUOffload(offload_params=True)

    # FSDP 配置
    model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        mixed_precision=mp_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,  # 完全分片
        cpu_offload=cpu_offload,
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,  # 预取优化
        limit_all_gathers=True,  # 限制 all-gather 防止 OOM
        use_orig_params=True,  # 保持原始参数（重要！）
    )

    return model

{
    "train_batch_size": 64,
    "gradient_accumulation_steps": 8,
    "train_micro_batch_size_per_gpu": 2,

    "bf16": {
        "enabled": true
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": 1e9,
        "stage3_prefetch_bucket_size": 1e8,
        "stage3_param_persistence_threshold": 1e6,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },

    "gradient_clipping": 1.0,

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 2e-5,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01
        }
    },

    "scheduler": {
        "type": "WarmupCosineLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 2e-5,
            "warmup_num_steps": 1000,
            "total_num_steps": 10000
        }
    },

    "activation_checkpointing": {
        "partition_activations": true,
        "cpu_checkpointing": false,
        "contiguous_memory_optimization": false,
        "number_checkpoints": null,
        "synchronize_checkpoint_boundary": false,
        "profile": false
    }
}

import time
import torch
from contextlib import contextmanager

@contextmanager
def profile_time(name):
    torch.cuda.synchronize()
    start = time.time()
    yield
    torch.cuda.synchronize()
    print(f"{name}: {time.time() - start:.2f}s")

def benchmark_training_step(model, batch, optimizer, use_fsdp=True):
    metrics = {}

    # Forward
    with profile_time("Forward"):
        output = model(batch)
        loss = output.loss
    metrics['loss'] = loss.item()

    # Backward
    with profile_time("Backward"):
        loss.backward()

    # Optimizer step
    with profile_time("Optimizer"):
        optimizer.step()
        optimizer.zero_grad()

    # 内存统计
    metrics['memory_allocated'] = torch.cuda.memory_allocated() / 1e9
    metrics['memory_reserved'] = torch.cuda.memory_reserved() / 1e9

    if use_fsdp:
        # FSDP 特定指标
        if hasattr(model, '_fsdp_wrapped_module'):
            metrics['communication_time'] = model._communication_time
    else:
        # DeepSpeed 特定指标
        if hasattr(optimizer, 'timer_names'):
            for name in optimizer.timer_names:
                metrics[f'ds_{name}'] = optimizer.timers(name)

    return metrics

# 迁移检查清单
migration_checklist = {
    "模型包装": "FSDP(...) -> deepspeed.initialize(...)",
    "优化器": "需要在 config 中配置，不能直接传入",
    "梯度累积": "自动处理 -> 需要显式配置",
    "Checkpoint": "torch.save -> model.save_checkpoint",
    "混合精度": "MixedPrecision -> fp16/bf16 config",
    "学习率调度": "手动 -> 配置文件",
}

# 迁移示例
def migrate_fsdp_to_deepspeed(fsdp_model, fsdp_optimizer):
    # 1. 提取原始模型
    if hasattr(fsdp_model, 'module'):
        base_model = fsdp_model.module
    else:
        base_model = fsdp_model

    # 2. 创建 DeepSpeed 配置
    ds_config = {
        "train_batch_size": world_size * batch_size,
        "gradient_accumulation_steps": grad_acc_steps,
        "zero_optimization": {
            "stage": 3 if was_full_shard else 2,
            # 其他配置...
        }
    }

    # 3. 初始化 DeepSpeed
    model, optimizer, _, _ = deepspeed.initialize(
        model=base_model,
        config=ds_config,
        model_parameters=base_model.parameters()
    )

    return model, optimizer

# FSDP 限制 all-gather
model = FSDP(
    model,
    limit_all_gathers=True,
    forward_prefetch=True,  # 前向预取
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
)

# 优化 ZeRO-3 配置
"zero_optimization": {
    "stage": 3,
    "stage3_max_live_parameters": 2e9,  # 增加缓存
    "stage3_max_reuse_distance": 2e9,
    "stage3_prefetch_bucket_size": 2e8,  # 增加预取
}

def convert_checkpoint(checkpoint_path, from_format="fsdp", to_format="deepspeed"):
    """转换不同格式的 checkpoint"""

    if from_format == "fsdp" and to_format == "deepspeed":
        # FSDP -> DeepSpeed
        state_dict = torch.load(checkpoint_path)

        # FSDP 可能有 _fsdp 前缀
        new_state_dict = {}
        for key, value in state_dict.items():
            new_key = key.replace("_fsdp_wrapped_module.", "")
            new_key = new_key.replace("_fpw_module.", "")
            new_state_dict[new_key] = value

        # DeepSpeed 期望的格式
        ds_checkpoint = {
            "module": new_state_dict,
            "epoch": 0,
            "global_step": 0,
        }

        torch.save(ds_checkpoint, checkpoint_path.replace(".pt", "_ds.pt"))

    elif from_format == "deepspeed" and to_format == "fsdp":
        # DeepSpeed -> FSDP
        # DeepSpeed ZeRO-3 需要先收集分片
        from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict

        convert_zero_checkpoint_to_fp32_state_dict(
            checkpoint_path,
            checkpoint_path.replace("ds", "fsdp.pt")
        )

def validation_step(rank, model, val_loader):
    if rank == 0:
        model.eval()
        total_loss = 0
        for batch in val_loader:
            loss = model(batch).loss
            total_loss += loss.item()
        avg_loss = total_loss / len(val_loader)
        torch.distributed.broadcast(torch.tensor([avg_loss]).cuda(), src=0)
    return avg_loss