File size: 4,941 Bytes
6011e08 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import logging
import time
import traceback
from pathlib import Path
import torch
from slime.utils.memory_utils import print_memory
logger = logging.getLogger(__name__)
class TrainProfiler:
def __init__(self, args):
self.args = args
self._torch_profiler_overall = None
self._memory_profiler_overall = None
if args.use_pytorch_profiler and ("train_overall" in args.profile_target):
self._torch_profiler_overall = _create_torch_profiler(args, name="train_overall")
if args.record_memory_history and ("train_overall" in args.profile_target):
self._memory_profiler_overall = _BaseMemoryProfiler.create(args)
self._memory_profiler_overall.start()
def on_init_end(self):
if self._torch_profiler_overall is not None:
self._torch_profiler_overall.start()
def step(self, rollout_id: int):
if self._torch_profiler_overall is not None:
self._torch_profiler_overall.step()
if (
self._memory_profiler_overall is not None
and ((s := self.args.memory_snapshot_num_steps) is not None)
and (rollout_id == s - 1)
):
self._memory_profiler_overall.stop()
def iterate_train_actor(self, iterator):
return _profile_simple_loop(iterator, self.args, name="train_actor")
def iterate_train_log_probs(self, iterator):
return _profile_simple_loop(iterator, self.args, name="train_log_probs")
def _profile_simple_loop(iterator, args, name):
if not (args.use_pytorch_profiler and (name in args.profile_target)):
yield from iterator
return
torch_profiler = _create_torch_profiler(args, name=name)
torch_profiler.start()
for item in iterator:
yield item
torch_profiler.step()
def _create_torch_profiler(args, name):
return torch.profiler.profile(
schedule=torch.profiler.schedule(
# TODO the train_actor and train_log_probs ones may need to have different args to control step
wait=max(args.profile_step_start - 1, 0),
warmup=1 if args.profile_step_start > 0 else 0,
active=args.profile_step_end - args.profile_step_start,
repeat=1,
),
on_trace_ready=torch.profiler.tensorboard_trace_handler(
args.tensorboard_dir,
worker_name=f"{name}_rank_{torch.distributed.get_rank()}",
use_gzip=True,
),
record_shapes=True,
with_stack=True,
profile_memory=True,
with_flops=True,
)
class _BaseMemoryProfiler:
@staticmethod
def create(args):
c = {
"torch": _TorchMemoryProfiler,
"memray": _MemrayMemoryProfiler,
}[args.memory_recorder]
return c(args)
def __init__(self, args):
self._path_dump = (
Path(args.memory_snapshot_dir)
/ f"memory_snapshot_time{time.time()}_rank{torch.distributed.get_rank()}_{args.memory_snapshot_path}"
)
def start(self):
raise NotImplementedError
def stop(self):
raise NotImplementedError
class _TorchMemoryProfiler(_BaseMemoryProfiler):
def start(self):
logger.info("Attach OOM dump memory history.")
torch.cuda.memory._record_memory_history(
max_entries=1000000,
# record stack information for the trace events
# trace_alloc_record_context=True,
stacks="all",
)
def oom_observer(device, alloc, device_alloc, device_free):
logger.info(
f"Observe OOM, will dump snapshot to {self._path_dump}. ({device=} {alloc=} {device_alloc=} {device_free=}; stacktrace is as follows)"
)
traceback.print_stack()
torch.cuda.memory._dump_snapshot(self._path_dump)
print_memory("when oom")
torch._C._cuda_attach_out_of_memory_observer(oom_observer)
def stop(self):
logger.info(f"Dump memory snapshot to: {self._path_dump}")
torch.cuda.memory._dump_snapshot(self._path_dump)
torch.cuda.memory._record_memory_history(enabled=None)
class _MemrayMemoryProfiler(_BaseMemoryProfiler):
def __init__(self, args):
super().__init__(args)
assert args.memory_snapshot_num_steps is not None, "In memray, must provide --memory-snapshot-num-steps"
def start(self):
logger.info("Memray tracker started.")
import memray
self._tracker = memray.Tracker(
file_name=self._path_dump,
native_traces=True,
)
self._tracker.__enter__()
def stop(self):
logger.info(f"Memray tracker stopped and dump snapshot to: {self._path_dump}")
self._tracker.__exit__(None, None, None)
|