feat: refactor main_ds.py (2/n) Accelerator class #594
Conversation
cdoern
force-pushed
the
refactor-accelerator
branch
from
1de1d23 to
25230a7
Compare
cdoern
force-pushed
the
refactor-accelerator
branch
2 times, most recently
from
aa7af32 to
81f731e
Compare
Accelerator works with Model to abstract common utilities behind a custom class that allows users to seamlessly setup their model for training Signed-off-by: Charlie Doern <[email protected]>
Signed-off-by: Charlie Doern <[email protected]>
cdoern
force-pushed
the
refactor-accelerator
branch
from
81f731e to
bf4be9f
Compare
|
E2E (NVIDIA L40S x4) (python 3.11) workflow launched on this PR: View run |
|
e2e workflow succeeded on this PR: View run, congrats! |
| class Accelerator: | ||
| def __init__( | ||
| self, | ||
| model: Model, |
There was a problem hiding this comment.
Model acts as a "factory" class that creates nn.module once the from_pretrained method is called in the case of each; Liger, Dolomite and the normal transformer model. That is what we should pass into the Accelerator so that we can avoid the weird model.model references.
There was a problem hiding this comment.
so, I think by not passing in model: Model we lose a lot of the seamless nature of these classes. things like self.model.lora_config are not possible within the Accelerator class if model is type hinted to nn.module, right?
There was a problem hiding this comment.
yeah, that's correct. I guess we need to refine the model class more for that to happen; approving rn in the spirit of getting the refactor in quickly.
thisisatharva-rh
left a comment
thisisatharva-rh
left a comment
There was a problem hiding this comment.
we need to have a larger conversation about what the Model class should look like, but for now, this looks good.
3 participants
** PLEASE NOTE, THIS PR INCLUDES THE CHANGES IN #572, AND WILL BE REDUCED IN SIZE ONCE THAT MERGES **
Introduce a new design for key components of main_ds.py. Namely splitting Model initialization, Accelerator initialization, Optimizer initialization, and Checkpoint saving initialization into classes. This commit introduces the Accelerator class
The Accelerator class aims to both store commonly accessed variables associated with the accelerated model and abstract model/optimizer mutation away from the user who should only access our Model and Optimizer classes.
These classes are one of a few steps needed to "SDK-ify" the training library
Adding structure to code via classes can either be someone's favorite or least favorite thing. So I figured I'd explain myself before continuing. Here is my rationale:
Classes provide logical structuring to code, especially code meant to be a publicly consumable SDK and allows you to associate related objects and methods with one another.
Being able to group functionality under the Model, Accelerator, and Checkpointer classes inherently reduces code complexity and duplication. Being able to store things like , self.distributed_framework,self.lora_config, etc in a way such that within the class they are accessible within different methods allows the arguments per method to go down drastically, as well as complex return values. Simpler methods and argument/return values allows for simpler testing of code.