xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-02 13:27:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Support gradient accumulation

Browse Source
This commit is contained in:
Kohaku-Blueleaf 2025-07-23 11:45:51 +08:00
parent 28a236d505
commit 2857b35703
1 changed files with 18 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
comfy_extras/nodes_train.py
View File

@ -39,13 +39,13 @@ def make_batch_extra_option_dict(d, indicies, full_size=None):
class TrainSampler(comfy.samplers.Sampler): class TrainSampler(comfy.samplers.Sampler):
def __init__(self, loss_fn, optimizer, loss_callback=None, batch_size=1, grad_acc=1, total_steps=1, seed=0, training_dtype=torch.bfloat16):
def __init__(self, loss_fn, optimizer, loss_callback=None, batch_size=1, total_steps=1, seed=0, training_dtype=torch.bfloat16):
self.loss_fn = loss_fn self.loss_fn = loss_fn
self.optimizer = optimizer self.optimizer = optimizer
self.loss_callback = loss_callback self.loss_callback = loss_callback
self.batch_size = batch_size self.batch_size = batch_size
self.total_steps = total_steps self.total_steps = total_steps
self.grad_acc = grad_acc
self.seed = seed self.seed = seed
self.training_dtype = training_dtype self.training_dtype = training_dtype
@ -92,6 +92,7 @@ class TrainSampler(comfy.samplers.Sampler):
self.loss_callback(loss.item()) self.loss_callback(loss.item())
pbar.set_postfix({"loss": f"{loss.item():.4f}"}) pbar.set_postfix({"loss": f"{loss.item():.4f}"})
if (i+1) % self.grad_acc == 0:
self.optimizer.step() self.optimizer.step()
self.optimizer.zero_grad() self.optimizer.zero_grad()
torch.cuda.empty_cache() torch.cuda.empty_cache()
@ -419,6 +420,16 @@ class TrainLoraNode:
"tooltip": "The batch size to use for training.", "tooltip": "The batch size to use for training.",
}, },
), ),
"grad_accumulation_steps": (
IO.INT,
{
"default": 1,
"min": 1,
"max": 1024,
"step": 1,
"tooltip": "The number of gradient accumulation steps to use for training.",
}
),
"steps": ( "steps": (
IO.INT, IO.INT,
{ {
@ -512,6 +523,7 @@ class TrainLoraNode:
positive, positive,
batch_size, batch_size,
steps, steps,
grad_accumulation_steps,
learning_rate, learning_rate,
rank, rank,
optimizer, optimizer,
@ -641,7 +653,8 @@ class TrainLoraNode:
optimizer, optimizer,
loss_callback=loss_callback, loss_callback=loss_callback,
batch_size=batch_size, batch_size=batch_size,
total_steps=steps, grad_acc=grad_accumulation_steps,
total_steps=steps*grad_accumulation_steps,
seed=seed, seed=seed,
training_dtype=dtype training_dtype=dtype
) )