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

Method _exec_optimizer_step

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

Source from the content-addressed store, hash-verified

1226 self.timers(PIPE_RECV_GRAD_TIMER).stop()
1227
1228 def _exec_optimizer_step(self, lr_kwargs=None):
1229 if self.wall_clock_breakdown():
1230 self.timers(STEP_MICRO_TIMER).start()
1231 self.timers(STEP_GLOBAL_TIMER).start()
1232
1233 self._force_grad_boundary = True
1234 self._take_model_step(lr_kwargs)
1235 self._force_grad_boundary = False
1236
1237 if self.global_rank == 0 and self.monitor.enabled:
1238 self.summary_events = [('Train/Samples/lr', self.get_lr()[0], self.global_samples)]
1239 loss_scale = self._get_optimizer_loss_scale() if self.fp16_enabled() else None
1240 if loss_scale is not None:
1241 self.summary_events.append(('Train/Samples/loss_scale', loss_scale, self.global_samples))
1242 self.monitor.write_events(self.summary_events)
1243
1244 if self.wall_clock_breakdown():
1245 self.timers(STEP_MICRO_TIMER).stop()
1246 self.timers(STEP_GLOBAL_TIMER).stop()
1247 if self.global_steps % self.steps_per_print() == 0:
1248 self.timers.log([
1249 BATCH_INPUT_TIMER,
1250 FORWARD_MICRO_TIMER,
1251 BACKWARD_MICRO_TIMER,
1252 BACKWARD_INNER_MICRO_TIMER,
1253 BACKWARD_REDUCE_MICRO_TIMER,
1254 STEP_MICRO_TIMER,
1255 ])
1256 if self.global_steps % self.steps_per_print() == 0:
1257 self.timers.log([
1258 FORWARD_GLOBAL_TIMER,
1259 BACKWARD_GLOBAL_TIMER,
1260 BACKWARD_INNER_GLOBAL_TIMER,
1261 BACKWARD_REDUCE_GLOBAL_TIMER,
1262 STEP_GLOBAL_TIMER,
1263 ])
1264
1265 def _allocate_zeros(self, shape, **kwargs):
1266 """ 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