One training step, in exact order

One training step, in exact order A sequence diagram generated by Archify. batch: input_ids, attention_mask, labels record every forward operation logits [batch, num_classes] loss: a single scalar loss.backward() gradients for every trainable parameter learning rate for this step gradients stored in .grad optimizer.step(), then zero_grad() Forward Backward Update DataLoader · collate + padding · Sequence participant DataLoader collate + padding Model · encoder + head · Sequence participant Model encoder + head Cross-entropy · logits vs labels · Sequence participant Cross-entropy logits vs labels Autograd · computation graph · Sequence participant Autograd computation graph AdamW · holds m_t and v_t · Sequence participant AdamW holds m_t and v_t Scheduler · per-step learning rate · Sequence participant Scheduler per-step learning rate Legend main request return security async trace message