[float] document e2e training -> inference flow

[danielvegamyhre]

Summary

Document the E2E training => inference flow with examples.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2190

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 723088a with merge base cdced21 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

cc @andrewor14

[andrewor14]

Tested this out locally too, works for me. Thanks!

[andrewor14]

(Might need to add a distributed checkpoint section but we can do that in a separate PR)

[danielvegamyhre]

In practice the model would be in some modeling file, and the training code and inference code would both import it separately, in order to avoid the need to deserialize the python model definition w/ torch.load(...., weights_only=False), which has some security risks.

However, I was aiming to have these be copy/paste-able runnable standalone examples, which seemed to require this bad practice. Thoughts @andrewor14 @vkuzo?

[vkuzo]

I think a good way to do it as as follows:

1. create a reproducible model definition
2. create a new instance of (1), train it, save weights to checkpoint
3. create a new instance of (1), load weights from checkpoint, finetune it or do inference

there is no saving of model definition with torch.save needed in the flow as above

[danielvegamyhre]

That's what I originally tried actually, but it doesn't work because the weights in the serialized/checkpointed model from step (2) are registered under different names (prefixed with _orig_mod) than the freshly initialized model in step (3).

I solved this by saving the converted model definition directly in torch.save and loading the model state dict into that, but it's not ideal imo. I'm curious how torchtitan/torchtune do this as well

[danielvegamyhre]

Ok going to merge this for now, we can discuss alternatives async if you want.