Spaces:
Running on Zero
Running on Zero
File size: 5,218 Bytes
9368ee7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from typing import TYPE_CHECKING
from lipforcing.callbacks.callback import Callback
from lipforcing.utils.distributed import world_size
import lipforcing.utils.logging_utils as logger
import torch
import wandb
try:
from torch.distributed.tensor import DTensor
except ImportError:
DTensor = None
if TYPE_CHECKING:
from lipforcing.methods import FastGenModel
def _get_local_numel(param: torch.Tensor) -> int:
"""Get the local (sharded) number of elements for a parameter.
For DTensor (FSDP2), returns the local shard size.
For regular tensors, returns the full size.
"""
if DTensor is not None and isinstance(param, DTensor):
return param._local_tensor.numel()
return param.numel()
class ParamCountCallback(Callback):
def on_train_begin(self, model: FastGenModel, **kwargs) -> None:
# get modules
modules = {"model": model, **model.model_dict}
# iterate over modules
output = {}
for name, module in modules.items():
# Logical (full model) param counts
trainable_params = sum(p.numel() for p in module.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in module.parameters())
# Local (sharded) param counts - what's actually in memory on this rank
local_trainable_params = sum(_get_local_numel(p) for p in module.parameters() if p.requires_grad)
local_total_params = sum(_get_local_numel(p) for p in module.parameters())
# check if parameter counts are different across ranks
if world_size() > 1:
trainable_params = self.gather_param_counts(trainable_params)
total_params = self.gather_param_counts(total_params)
local_trainable_params = self.gather_param_counts(local_trainable_params)
local_total_params = self.gather_param_counts(local_total_params)
if len(set(total_params)) == 1 and len(set(trainable_params)) == 1:
trainable_params = trainable_params[0]
total_params = total_params[0]
if len(set(local_total_params)) == 1 and len(set(local_trainable_params)) == 1:
local_trainable_params = local_trainable_params[0]
local_total_params = local_total_params[0]
# logging
module_name = module.__class__.__name__
output.update(
{
f"{name}/trainable_params": trainable_params,
f"{name}/total_params": total_params,
f"{name}/local_trainable_params": local_trainable_params,
f"{name}/local_total_params": local_total_params,
}
)
if isinstance(trainable_params, list):
logger.warning(f"Parameter counts differ across ranks for {module_name}.")
for rank, (p_train, p) in enumerate(zip(trainable_params, total_params)):
logger.info(
f"{name} ({module_name}) has {p_train * 1.e-6:.2f} M trainable and {p * 1.e-6:.2f} M total params on rank {rank}."
)
else:
logger.info(
f"{name} ({module_name}) has {trainable_params * 1.e-6:.2f} M trainable and {total_params * 1.e-6:.2f} M total params (logical)."
)
# Report local/sharded counts
if isinstance(local_trainable_params, list):
for rank, (p_train, p) in enumerate(zip(local_trainable_params, local_total_params)):
logger.info(
f"{name} ({module_name}) has {p_train * 1.e-6:.2f} M trainable and {p * 1.e-6:.2f} M total params LOCAL on rank {rank}."
)
else:
is_sharded = local_total_params < total_params if not isinstance(total_params, list) else True
if is_sharded:
logger.info(
f"{name} ({module_name}) has {local_trainable_params * 1.e-6:.2f} M trainable and {local_total_params * 1.e-6:.2f} M total params LOCAL per rank (sharding ratio: {world_size()}x)."
)
else:
logger.info(f"{name} ({module_name}) is NOT sharded (local == logical params).")
if wandb.run:
wandb.run.summary.update(output)
def gather_param_counts(self, param_count):
"""
Gather parameter counts across all ranks.
Args:
param_count: Parameter count to gather.
Returns:
List of parameter counts across all ranks.
"""
param_count = torch.tensor(
[param_count], dtype=torch.long, device="cuda" if torch.cuda.is_available() else "cpu"
)
param_count_list = [torch.zeros_like(param_count) for _ in range(world_size())]
torch.distributed.all_gather(param_count_list, param_count)
return [p.item() for p in param_count_list]
|