Spaces:
Running on Zero
Running on Zero
File size: 3,328 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
from typing import TYPE_CHECKING, Callable, Any, Dict, List
import pandas as pd
import psutil
import torch
from lipforcing.callbacks.callback import Callback
from lipforcing.utils.distributed import world_size, is_rank0, synchronize
import lipforcing.utils.logging_utils as logger
if TYPE_CHECKING:
from lipforcing.methods import FastGenModel
def log_prof_data(data_list: List[Dict[str, Any]]):
# Create a table to log data with rank information
metrics = list(data_list[0].keys())
# Initialize dictionaries to store min and max values for each metric
min_values = {key: float("inf") for key in metrics}
max_values = {key: float("-inf") for key in metrics}
sum_values = {key: 0.0 for key in metrics}
count = 0
for _rank, prof_data in enumerate(data_list):
count += 1
# Update min, max, and sum values
for key in metrics:
min_values[key] = min(min_values[key], prof_data[key])
max_values[key] = max(max_values[key], prof_data[key])
sum_values[key] += prof_data[key]
# Calculate average values
avg_values = {key: sum_values[key] / count for key in metrics}
summary_df = pd.DataFrame({"Avg": avg_values, "Max": max_values, "Min": min_values})
logger.info(f"GPU stats:\n{summary_df.to_string()}")
class GPUStatsCallback(Callback):
def __init__(self, every_n: int = 100):
self.every_n = every_n
def on_train_begin(self, model: FastGenModel, iteration: int = 0):
torch.cuda.reset_peak_memory_stats()
if hasattr(self, "config"):
# overwritten by logging_iter if self.config exists
self.every_n = self.config.trainer.logging_iter
logger.info(f"every_n to measure gpus stats: {self.every_n}")
def on_training_step_end(
self,
model: FastGenModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor | Callable],
loss_dict: dict[str, torch.Tensor],
iteration: int = 0,
) -> None:
del data_batch, output_batch, loss_dict
if iteration % self.every_n == 0:
cur_process = psutil.Process(os.getpid())
cpu_memory_usage = sum(p.memory_info().rss for p in [cur_process] + cur_process.children(recursive=True))
cpu_mem_gb = cpu_memory_usage / (1024**3)
peak_gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
peak_gpu_mem_reserved_gb = torch.cuda.max_memory_reserved() / (1024**3)
util = torch.cuda.utilization()
prof_data = {
"cpu_mem_gb": float(cpu_mem_gb),
"peak_gpu_mem_gb": float(peak_gpu_mem_gb),
"peak_gpu_mem_reserved_gb": float(peak_gpu_mem_reserved_gb),
"util": float(util),
}
synchronize()
data_list = [prof_data] * world_size()
# this is blocking by default
if world_size() > 1:
torch.distributed.all_gather_object(data_list, prof_data)
if is_rank0():
log_prof_data(data_list)
synchronize()
|