[feature request] Support for multiple models

[Cameron7195]

For my use case I would find it extremely useful to perform distillation with a KL objective, with a larger teacher model (gradient-free) and smaller student model which has losses/gradients.

I found this feature implemented in deepspeed last year:
huggingface/accelerate#2496
huggingface/accelerate#3097

https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed_multiple_model
So, it's possible that doing multiple models is just a matter of following the tutorial^ and applying it to model loading in the trainer.__init__:

ArcticTraining/arctic_training/trainer/trainer.py

Line 148 in 2544f03

def __init__(self, config: TrainerConfig, mode: str = "train") -> None:

if this is the case that would be awesome- and probably means I can accomplish this solely through creating a custom Trainer class, which I understand is expected of users anyway.

Reading through this tutorial / issues + the arctic_training.trainer, I think what is needed is to rewrite model loading & engine initialization using utils.DeepSpeedPlugins. My main confusion is I don't know where the accelerator init and prepare calls are happening in ArcticTraining- or whether they happen at all. How can the deepspeed/accelerate setup be mapped onto the multiple models tutorial with this in mind? The accelerator calls seem necessary to follow the exact flow in the tutorial.

tldr; this feature might be purely expected of users to write, or may require some rejigging of the trainer, I am just not sure.

[sfc-gh-mwyatt]

Hi @Cameron7195 thanks for the issue! We don't rely on accelerate to do the init of deepspeed models. Instead, we're calling deepspeed.initialize directly:

ArcticTraining/arctic_training/trainer/trainer.py

Line 228 in 8f61c50

self.model, *_ = deepspeed.initialize(

However, what you describe is still totally doable with a custom trainer and/or custom model factory. For example, in the mlp speculator project we have a custom model factory that loads a secondary speculator model:

ArcticTraining/projects/mlp_speculator/train.py

Line 84 in 8f61c50

class MLPSpeculatorModelFactory(HFModelFactory):

Ideally we want to keep these specific types of changes outside the base Trainer class and instead implement them in class that inherits from Trainer. This keeps the base class a bit cleaner and helps avoid breaking changes for other user-implemented trainers.

[Cameron7195]

This worked! Thanks so much for the help