File size: 6,038 Bytes
b7eca4a | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | # Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import datetime
import inspect
import logging
from contextlib import contextmanager
from typing import Any, Dict, Tuple
import torch
import torch.distributed as dist
from codetiming import Timer
from verl.utils.device import get_torch_device
from verl.utils.logger.aggregate_logger import DecoratorLoggerBase
def _get_current_mem_info(unit: str = "GB", precision: int = 2) -> Tuple[str]:
"""Get current memory usage."""
assert unit in ["GB", "MB", "KB"]
divisor = 1024**3 if unit == "GB" else 1024**2 if unit == "MB" else 1024
mem_allocated = get_torch_device().memory_allocated()
mem_reserved = get_torch_device().memory_reserved()
# use get_torch_device().mem_get_info to profile device memory
# since vllm's sleep mode works below pytorch
# see https://github.com/vllm-project/vllm/pull/11743#issuecomment-2754338119
mem_free, mem_total = get_torch_device().mem_get_info()
mem_used = mem_total - mem_free
mem_allocated = f"{mem_allocated / divisor:.{precision}f}"
mem_reserved = f"{mem_reserved / divisor:.{precision}f}"
mem_used = f"{mem_used / divisor:.{precision}f}"
mem_total = f"{mem_total / divisor:.{precision}f}"
return mem_allocated, mem_reserved, mem_used, mem_total
def log_gpu_memory_usage(head: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0):
if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank):
mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info()
message = f"{head}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, device memory used/total (GB): {mem_used}/{mem_total}"
if logger is None:
print(message)
else:
logger.log(msg=message, level=level)
class GPUMemoryLogger(DecoratorLoggerBase):
"""A decorator class to log GPU memory usage.
Example:
>>> from verl.utils.debug.performance import GPUMemoryLogger
>>> @GPUMemoryLogger(role="actor")
>>> def update_actor(self, batch):
... # real actor update logics
... return
"""
def __init__(self, role: str, logger: logging.Logger = None, level=logging.DEBUG, log_only_rank_0: bool = True):
if dist.is_initialized() and dist.get_world_size() > 1:
rank = dist.get_rank()
else:
rank = 0
super().__init__(role, logger, level, rank, log_only_rank_0)
def __call__(self, decorated_function: callable):
def f(*args, **kwargs):
return self.log(decorated_function, *args, **kwargs)
return f
def log(self, func, *args, **kwargs):
name = func.__name__
mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info()
message = f"Before {name}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, device memory used/total (GB): {mem_used}/{mem_total}"
self.logging_function(message)
output = func(*args, **kwargs)
mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info()
message = f"After {name}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, device memory used/total (GB): {mem_used}/{mem_total}"
self.logging_function(message)
return output
def log_print(ctn: Any):
current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
frame = inspect.currentframe().f_back
function_name = frame.f_code.co_name
line_number = frame.f_lineno
file_name = frame.f_code.co_filename.split("/")[-1]
print(f"[{current_time}-{file_name}:{line_number}:{function_name}]: {ctn}")
@contextmanager
def _timer(name: str, timing_raw: Dict[str, float]):
"""Context manager for timing code execution.
This utility function measures the execution time of code within its context
and accumulates the timing information in the provided dictionary.
Args:
name (str): The name/identifier for this timing measurement.
timing_raw (Dict[str, float]): Dictionary to store timing information.
Yields:
None: This is a context manager that yields control back to the code block.
"""
with Timer(name=name, text="{name}: {seconds:.1f} seconds", logger=None) as timer:
yield
if name not in timing_raw:
timing_raw[name] = 0
timing_raw[name] += timer.last
def reduce_timing(timing_raw: Dict[str, float]) -> Dict[str, float]:
"""Reduce timing information across all processes.
This function uses distributed communication to gather and sum the timing
information from all processes in a distributed environment.
Args:
timing_raw (Dict[str, float]): Dictionary containing timing information.
Returns:
Dict[str, float]: Reduced timing information.
"""
if not dist.is_initialized():
return timing_raw
key_list, timing_list = [], []
for key in sorted(timing_raw.keys()):
key_list.append(key)
timing_list.append(timing_raw[key])
timing_list = torch.tensor(timing_list, dtype=torch.float32, device=torch.cuda.current_device())
torch.distributed.all_reduce(timing_list, op=torch.distributed.ReduceOp.AVG)
timing_list = [tensor.item() for tensor in timing_list.to("cpu")]
timing_generate = {key_list[i]: timing_list[i] for i in range(len(key_list))}
return timing_generate
|