Runtime Performance Evaluation and NVTX Profiling
This page explains how to use the optim_test_tools to measure runtime and add NVTX ranges for profiling.
We focus on the most important measurements for the training iteration and loss/optimizer computation.
We rely on two singletons from accvlab.optim_test_tools:
Stopwatch: lightweight runtime measurements with optional CPU usage samplingNVTXRangeWrapper: NVTX push/pop helpers with additional functionality to annotate critical sections for Nsight Systems (& Nsight Compute).
Enable & Configure in the Training Script
The configuration is centralized in the main training script (demos/streampetr/tools/train.py). The
runtime measurement and NVTX range wrapper are enabled based on the runtime_meas_config section in the
config. The use of the configuration file enables us to use different configurations for different types of
runs, e.g.:
Training (actual usage): All disabled
Evaluation: Only Stopwatch enabled
Profiling: All enabled, nvtx range mismatches not checked (to avoid overhead)
(Debugging: All enabled, nvtx range mismatches checked (to help debugging range mismatches))
Note that the Stopwatch and NVTXRangeWrapper are enabled & configured centrally here, and can be used
throughout the codebase relying on this configuration.
281 # ===== Runtime Measurement & Profiling Configuration =====
282 #@NOTE
283 # Whether to use the runtime measurements and the nvtx ranges is configurable. Here, we read the
284 # configuration from the config file and enable the corresponding signletons if needed.
285 runtime_meas_cfg = cfg.get("runtime_meas_config", None)
286 if runtime_meas_cfg is not None:
287 stopwatch_cfg = runtime_meas_cfg.stopwatch
288 nvtx_rng_wrp_cfg = runtime_meas_cfg.nvtx_range_wrapper
289 if stopwatch_cfg.enable and rank == 0:
290 stopwatch.enable(stopwatch_cfg.num_warmup_iters, stopwatch_cfg.print_every_n_iters, stopwatch_cfg.do_cuda_sync)
291 if nvtx_rng_wrp_cfg.enable:
292 nvtx_wrp.enable(nvtx_rng_wrp_cfg.sync_on_push, nvtx_rng_wrp_cfg.sync_on_pop, nvtx_rng_wrp_cfg.keep_track_of_range_order)
293 # ===== End of Runtime Measurement & Profiling Configuration =====
Measuring the Training Iteration
The core iteration timings are implemented in the custom runner
(IterBasedRunnerWithTime) to capture the outer training loop and the inner training step,
including optional CPU usage measurement. Note that this means that this runner needs to be used instead
of the original implementation (from the mmcv package) to enable the runtime measurements. The selection
can be done by setting runner.type in the config file to IterBasedRunnerWithTime.
Outer loop timing with CPU usage sampling in run():
184 # ===== Training loop & timing =====
185
186 while self.iter < self._max_iters:
187 #@NOTE
188 # In this loop, we use the stopwatch to
189 # - Time the training loop
190 # - Measure the CPU usage of the training loop
191 # - Indicate the end of an iteration
192 # Also see `train()` method in this class for further stopwatch usage.
193
194 #@NOTE
195 # A single measurement can be configured to also measure the CPU usage. We want to measure
196 # the CPU usage for the training loop, and set it up here by (referenced by the name of the
197 # measurement).
198 stopwatch.set_cpu_usage_meas_name("training loop")
199 for i, flow in enumerate(workflow):
200 #@NOTE
201 # We start the measurement for the training loop (including CPU usage measurement, see above).
202 #
203 # Note that in the loop, the `train()` method is called internally.
204 stopwatch.start_meas("training loop")
205 self._inner_iter = 0
206 mode, iters = flow
207 if not isinstance(mode, str) or not hasattr(self, mode):
208 raise ValueError(
209 'runner has no method named "{}" to run a workflow'.
210 format(mode))
211 iter_runner = getattr(self, mode)
212 for _ in range(iters):
213 if mode == 'train' and self.iter >= self._max_iters:
214 break
215 iter_runner(iter_loaders[i], **kwargs)
216 #@NOTE: We end the measurement for the training loop.
217 stopwatch.end_meas("training loop")
218 #@NOTE
219 # Indicate the end of an iteration. This is important, as this allows the Stopwatch to
220 # - Average the runtime correctly for individual iterations
221 # - Decide when to end the warm-up phase and start the measurements (after warm-up phase)
222 # - Decide when to auto-print the measurements (every n iterations)
223 stopwatch.finish_iter()
224
225 # ===== End of Training loop & timing =====
Training step timing in train():
84 # ===== Training function & timing =====
85
86 def train(self, data_loader, **kwargs):
87
88 #@NOTE
89 # We use the stopwatch throughout this method to time different parts of the training loop.
90 # Also see `run()` method for further stopwatch usage (measuring CPU usage,
91 # indicating the end of an iteration).
92
93 #@NOTE: Start the measurement for the training iteration.
94 stopwatch.start_meas("IterBasedRunnerWithTime.train")
95 self.model.train()
96 self.mode = 'train'
97 self.data_loader = data_loader
98 self._epoch = data_loader.epoch
99 data_batch = next(data_loader)
100 self.data_batch = data_batch
101 self.call_hook('before_train_iter')
102 #@NOTE: Start the measurement for the training forward step.
103 stopwatch.start_meas("Training step")
104 outputs = self.model.train_step(data_batch, self.optimizer, **kwargs)
105 #@NOTE: End the measurement for the training forward step.
106 stopwatch.end_meas("Training step")
107 if not isinstance(outputs, dict):
108 raise TypeError('model.train_step() must return a dict')
109 if 'log_vars' in outputs:
110 self.log_buffer.update(outputs['log_vars'], outputs['num_samples'])
111 self.outputs = outputs
112 self.call_hook('after_train_iter')
113 del self.data_batch
114 self._inner_iter += 1
115 self._iter += 1
116 #@NOTE: End the measurement for the training iteration.
117 stopwatch.end_meas("IterBasedRunnerWithTime.train")
118
119 #@NOTE
120 # Note that this function does not explicitly call the optimizer step. Instead, the optimizer step is
121 # added as a hook, and so is called implicitly here. Further runtime measurements & nvtx ranges are
122 # added to the optimizer (see `optimizer_with_time.py` in the same directory as this file).
123
124 # ===== End of Training function & timing =====
Measuring Optimizer/Loss Step and NVTX Ranges
The optimizer step is instrumented with both Stopwatch and NVTXRangeWrapper to
annotate and time the critical loss backward/step section. Note that similar to the runner, this
implementation needs to be used instead of the original implementation by setting up the optimizer_config
parameter in the config file with type='Fp16OptimizerHookWithTime'.
153 # ===== Optimizer Step Timing and NVTX Ranges =====
154 # clear grads of last iteration
155 #@NOTE
156 # Start runtime measurement & push NVTX range.
157 nvtx_wrp.range_push("optimizer.after_train_iter")
158 stopwatch.start_meas("Optimizer step")
159
160 runner.model.zero_grad()
161 runner.optimizer.zero_grad()
162
163 self.loss_scaler.scale(runner.outputs['loss']).backward()
164 self.loss_scaler.unscale_(runner.optimizer)
165 # grad clip
166 if self.grad_clip is not None:
167 grad_norm = self.clip_grads(runner.model.parameters())
168 if grad_norm is not None:
169 # Add grad norm to the logger
170 runner.log_buffer.update({'grad_norm': float(grad_norm)},
171 runner.outputs['num_samples'])
172 # backward and update scaler
173 self.loss_scaler.step(runner.optimizer)
174 self.loss_scaler.update(self._scale_update_param)
175
176 # save state_dict of loss_scaler
177 runner.meta.setdefault(
178 'fp16', {})['loss_scaler'] = self.loss_scaler.state_dict()
179
180 #@NOTE
181 # Pop NVTX range & end runtime measurement.
182 nvtx_wrp.range_pop()
183 stopwatch.end_meas("Optimizer step")
184 # ===== End of Optimizer Step Timing and NVTX Ranges =====
NVTX ranges also identify operations in the matching implementation and the loss computations described in the Batching Helpers demo.
NVTX Range Mismatch Checking
To enable NVTX range mismatch checking, set both enable and keep_track_of_range_order to True
in runtime_meas_config.nvtx_range_wrapper. When checking is enabled, calls to
range_pop(expected_name) verify that the range being closed has the expected name. Disable
keep_track_of_range_order during performance profiling to avoid the checking overhead.
For example, HungarianAssigner2DBatched checks its outermost range when the assignment finishes:
320 #@NOTE: Types of range pops
321 #@NOTE: A simple range pop to close the NVTX range, without mismatch checking.
322 nvtx_wrp.range_pop()
323 #@NOTE: This will check for range stack mismatches if checking is enabled (enable only for debugging).
324 nvtx_wrp.range_pop("HungarianAssigner2DBatched.assign")
Viewing Results
Stopwatchresults are printed from the main script at the end of training:demos/streampetr/tools/train.py521 # ===== Stopwatch: Print eval times ===== 522 #@NOTE 523 # Print the run-time measurements at the end of the training. Note that depending on the setup, the 524 # measurements will also be printed periodically (e.g. every 50 iterations). Here, we ensure that the 525 # last printed measurement is for the complete training. 526 stopwatch.print_eval_times() 527 # ===== End of Stopwatch: Print eval times =====
For NVTX, launch your run under Nsight Systems/Compute. The training script additionally uses
torch.cuda.cudart().cudaProfilerStart()andtorch.cuda.cudart().cudaProfilerStop()to start and stop the profiling session. Note that this is done only if the use of NVTX is enabled in the config (and therefore, the NVTX range wrapper is enabled).demos/streampetr/tools/train.py486 # ===== Start profiling (if running with nsys) ===== 487 #@NOTE: Indicate that profiling should be started if the NVTX range wrapper is enabled. 488 if nvtx_wrp.is_enabled: 489 if "RANK" in os.environ and int(os.environ["RANK"]) == 0: 490 print("Starting profiling") 491 torch.cuda.cudart().cudaProfilerStart() 492 # ===== End of Start profiling (if running with nsys) ===== 493 494 495 # ===== Train the model ===== 496 #@NOTE 497 # Call the training function. Note that datasets is passed to the model. As this now may be a DALI 498 # iterator, one needs to make sure that the training function is able to handle this (by using it directly 499 # as a drop-in replacement for a PyTorch DataLoader, instead of creating a PyTorch DataLoader from it). 500 # This is done in `projects/mmdet3d_plugin/core/apis/mmdet_train.py`. 501 custom_train_model( 502 model, 503 datasets, 504 cfg, 505 distributed=distributed, 506 validate=(not args.no_validate), 507 timestamp=timestamp, 508 meta=meta) 509 # ===== End of Train the model ===== 510 511 512 # ===== Stop profiling (if running with nsys) ===== 513 #@NOTE: Indicate that profiling should be stopped if the NVTX range wrapper is enabled. 514 if nvtx_wrp.is_enabled: 515 if "RANK" in os.environ and int(os.environ["RANK"]) == 0: 516 print("Stopping profiling") 517 torch.cuda.cudart().cudaProfilerStop() 518 # ===== End of Stop profiling (if running with nsys) =====