Nodes/ComfyDL/SGD Step
ComfyUI Node

SGD Step

The node that actually trains your model

By Cynthia-lxx·Created 2 months ago·Updated 2 days ago· 6
SGD Step
  • model
  • model
lr0.030
batch_size32

Everything before this node is plumbing - forward pass, loss, gradients. CdlSgdStep is where the model actually changes. It applies one mini-batch stochastic gradient descent update: every parameter moves a little down its gradient, scaled by the learning rate and divided by the batch size, then the gradients are zeroed for the next iteration. In a ComfyDL training loop, this is the node that turns your graph into a learner.

It's the sgd function from the Dive into Deep Learning book, and it's deliberately minimal - no momentum, no Adam, no fancy schedules. Just textbook SGD, the version that appears on page one of every optimizer chapter, so you can see exactly what's happening before anyone dresses it up.

The inputs

  • lr - learning rate (default 0.03). The one knob that matters most. Too big and the loss diverges; too small and training is glacially slow. The d2l examples live in the 0.01–0.1 range.
  • batch_size - used as the normalizer in lr * grad / batch_size (default 32). This is the mini-batch size your data-loader used; keep them consistent.
  • model - the cdlModel to update (optional; nothing happens if you leave it unwired, and you get a None back).

Output: model - the same model, now with updated weights. You feed it back into the next forward pass to continue the loop.

The critical prerequisite

Gradients must already exist. SGD Step assumes a loss.backward() has run on your model's parameters - that's what fills in param.grad. In ComfyDL, gradient computation happens inside a loss/backward node before this one in the graph. If you wire SGD Step in too early, there's nothing to step on; the node just passes the model through unchanged. The execution order of the graph is doing the work that optimizer.step() normally does in a script.

Also note it zeros gradients after the update (param.grad.zero_()), which is exactly what you want - accumulate gradients over a batch, step once, clean up.

Installing ComfyDL

Standard light install:

cd ComfyUI/custom_nodes
git clone https://github.com/Cynthia-lxx/ComfyDL
cd ComfyDL && pip install -r requirements.txt

Restart ComfyUI; it's under ComfyDL → TorchOps, or search "ComfyDL" in ComfyUI Manager. Only extra dependency is matplotlib; no downloads.

Building the loop

The classic ComfyDL training cycle in nodes: synthetic or dataset batch → forward pass → loss → backward → SGD Step → (repeat, feeding the updated model back around). Since ComfyUI isn't a loop construct, d2l workflows here tend to unroll a handful of steps, or you re-run the graph and watch the loss curve in Plot descend. That's the whole "training a network in a node graph" experience, and it's surprisingly legible.

The one trap: don't crank lr past ~0.1 on a small dataset and expect stability - you'll watch the loss go up, which is at least visually educational. And keep batch_size honest: it's a fixed normalizer, not an auto-detected value, so if you change your data batching, change this to match or your step sizes drift.

Categoryd2l/TorchOps

Inputs (3)

NameTypeDefaultDescription
lrFLOAT0.0301e-8–10
batch_sizeINT321–65536
modeloptcdlModel

Outputs (1)

NameTypeDescription
modelcdlModel