Train Medusa Demo. via: https://huggingface.co/docs/text-generation-inference/basic_tutorials/train_medusa