Significant degradation between versions

[keveman]

I am attaching a graph comparing loss curves for my model across two commits of HeavyBall be36047 and c48bcaf There was no other change in my code, not in the model or data. My call to the optimizer looks as follows:

optimizer = heavyball.ForeachPSGDKron(
        opt_grouped_params,
        lr=2e-3,
        weight_decay=5e-4,
        warmup_steps=1024*16,
)

Granted c48bcaf is not current HEAD, but I was taken aback by this sudden and drastic loss of performance. @ClashLuke Any hints?

[keveman]

@ClashLuke This is mostly a fact finding exercise for me and I don't want you to treat this as a bug. I'm yet to try HEAD on my model, but waiting for some GPU resources.

[ClashLuke]

Thank you for the detailed report! I don't immediately see what went wrong, but I'll get back to you.

The most recent change was using a more accurate spectral norm calculation, which changes the effective precond_lr

[ClashLuke]

Could you try setting precond_init_scale=1?

[keveman]

Thanks for looking into this @ClashLuke. I am not sure whether you meant trying precond_init_scale=1 at c48bcaf or HEAD, but I'll try both and report back as soon as I get some GPU time.

[ClashLuke]

HEAD and c48bcaf don't a have any significant differences for PSGD, so, it's definitely an unfixed bug. Hopefully we can fix it soon

[keveman]

Sorry it took this long, but here is the comparison with precond_init_scale=1.

No change pretty much with precond_init_scale=1, so definitely something else changed after be36047 causing the dramatic difference.

[ClashLuke]

Could you share any information about the problem to help me reproduce the curves? (or modify one of the example files to match your setup)

Otherwise, the other major changes can be disabled using

1. dampening=0
2. lower_bound_beta=0
3. precond_update_power_iterations=1

The updated PSGD also balances the preconditioners at every step, slightly changing learning dynamics. This didn't break anything in my testing, but you could comment the psgd_balance_Q lines here:

HeavyBall/heavyball/utils.py

Lines 2025 to 2031 in c48bcaf

	@decorator_knowngood
	def _balance_to_triu(Q: "TriuOrLine"):
	if isinstance(Q[0], tuple):
	psgd_balance_Q([o[1] for o in Q])
	return line_to_triu(Q)
	psgd_balance_Q(Q)
	return Q

[keveman]

Thanks for looking into this @ClashLuke. The model I running is here https://huggingface.co/UsefulSensors/moonshine-tiny
I am running on 64 GPUs with the accelerate framework. I'll try to repro this with single GPU setting, and also try modify the example files like you suggested, and report back.

[ClashLuke]

Can't reproduce a regression with autoencoder.py across various model and batch sizes, as well as with and without the flags above and independently testing cache (the only other thing that was changed) - the curves are almost identical.

I did notice that stochastic_schedule was changed from True to False and that the corresponding argument was ignored, which is now fixed

The only noteworthy things seem to be the flags mentioned above

[anhnguyen-namitech]

Had a similar result. ForeachAdamW converges so slowly compared to AdamW on a speech-to-text task.

[ClashLuke]

Sorry, could you clarify your message above @anhnami ? Does heavyball.ForeachAdamW converge slower than torch.optim.AdamW? If so, please make sure the hyperparameters match.

This is not documented, but heavyball has different defaults than torch.optim, optax and others. heavyball's implementation is aligned closely with that of torch, except for slight precision improvements. However, its beta2 is set to 0.99 compared to torch's 0.999.

[anhnguyen-namitech]

Hi, sorry for the late reply. In May, we trained Conformer RNN-T ASR model with heavyball.ForeachAdamW and torch.optim.AdamW using identical settings: lr=1e-3, wd=1e-3, betas[0.9, 0.98], batch size: 2048. For some unknown reasons, ForeachAdamW converged much slower than AdamW (similar to training curves showed by keveman). Unfortunately, I have removed the training logs for them.

[ClashLuke]

Hey @keveman, I might've finally made some progress on this!

From the 2.1.4 release docs:

psgd's scheduler was not stepped properly, causing the scheduler to remain at 100% precond update probability

Would love to know if this fixes it for you

[AlexDeMoor]

Hi @ClashLuke , I m encountering a similar issue for my use case. Everything works well with the old version we used (1.3 or 1.4) however since we switched to the early 2.0's or the most recent one we see a degradation in the convergence of our models (applications are in particle physics). Our setup was quite simple, using ForeachPSGDKron with beta=0.7, precond_lr=0.05. We tried implementing dampening=0, lower_bound_beta=0, precond_update_power_iterations=1 as suggested above but this also didn't help us.