Okay, thanks for the details. I have access to a dual-A100 machine for now, so I'm going to capitalize on the opportunity. I wound up writing my own lightning-based training code, as I wanted to dig in a bit more, to get used to how lightning does things. I dig it! The modularity is a welcome change from vanilla and/or accelerate scripts.

I've been working from the Mousai paper, but just realized that I probably shouldn't have added a cosine scheduler—the paper says nothing about a lr schedule, so I'm guessing it was just 1e-4 throughout. I'm going to try resuming with that setting now... fingers crossed.

I also don't currently use EMA, which I should really add...