MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / _exec_optimizer_step

Method _exec_optimizer_step

deepspeed/runtime/pipe/engine.py:1244–1279  ·  view source on GitHub ↗
(self, lr_kwargs=None)

Source from the content-addressed store, hash-verified

1242 self.timers(PIPE_RECV_GRAD_TIMER).stop()
1243
1244 def _exec_optimizer_step(self, lr_kwargs=None):
1245 if self.wall_clock_breakdown():
1246 self.timers(STEP_MICRO_TIMER).start()
1247 self.timers(STEP_GLOBAL_TIMER).start()
1248
1249 self._force_grad_boundary = True
1250 self._take_model_step(lr_kwargs)
1251 self._force_grad_boundary = False
1252
1253 if self.global_rank == 0 and self.monitor.enabled:
1254 self.summary_events = [('Train/Samples/lr', self.get_lr()[0], self.global_samples)]
1255 loss_scale = self._get_optimizer_loss_scale() if self.fp16_enabled() else None
1256 if loss_scale is not None:
1257 self.summary_events.append(('Train/Samples/loss_scale', loss_scale, self.global_samples))
1258 self.monitor.write_events(self.summary_events)
1259
1260 if self.wall_clock_breakdown():
1261 self.timers(STEP_MICRO_TIMER).stop()
1262 self.timers(STEP_GLOBAL_TIMER).stop()
1263 if self.global_steps % self.steps_per_print() == 0:
1264 self.timers.log([
1265 BATCH_INPUT_TIMER,
1266 FORWARD_MICRO_TIMER,
1267 BACKWARD_MICRO_TIMER,
1268 BACKWARD_INNER_MICRO_TIMER,
1269 BACKWARD_REDUCE_MICRO_TIMER,
1270 STEP_MICRO_TIMER,
1271 ])
1272 if self.global_steps % self.steps_per_print() == 0:
1273 self.timers.log([
1274 FORWARD_GLOBAL_TIMER,
1275 BACKWARD_GLOBAL_TIMER,
1276 BACKWARD_INNER_GLOBAL_TIMER,
1277 BACKWARD_REDUCE_GLOBAL_TIMER,
1278 STEP_GLOBAL_TIMER,
1279 ])
1280
1281 def _allocate_zeros(self, shape, **kwargs):
1282 """ Allocate a tensor of zeros on the engine's device.

Callers

nothing calls this directly

Calls 11

wall_clock_breakdownMethod · 0.80
_take_model_stepMethod · 0.80
appendMethod · 0.80
steps_per_printMethod · 0.80
startMethod · 0.45
get_lrMethod · 0.45
fp16_enabledMethod · 0.45
write_eventsMethod · 0.45
stopMethod · 0.45
logMethod · 0.45

Tested by

no test coverage detected