implement post training

[ehartford]

Does this look right?

[casper-hansen]

So the purpose of post_training is to recreate the original module with the trained weights. Ideally, you would add an original_class variable to the __init__ and save that for later. Then once you are done training, post_training is called to make your checkpoint compatible with the original model - so the idea here would be that original_class=MixtralSparseMoeBlock.

[ehartford]

Is this closer?

[winglian]

Can you provide more details please about what this solves? Sorry that I'm not understanding from the provided changes et.

[casper-hansen]

@winglian this is supposed to solve the conversion back from the fused MoE into the original architecture.

[ehartford]

I am trying to help get the branch to a usable state because I wanna start training a model with it