[loading] Fix forced upcasting to fp32 - #43683
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: afmoe |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Looks overall good to me just one thing to remove (dtype plan is still in the signature of the core model loadin fn) + 2 small questions
| self.shared_experts = AfmoeMLP(config, config.moe_intermediate_size * config.num_shared_experts) | ||
| self.experts = AfmoeExperts(config) | ||
| self.expert_bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False) | ||
| self.expert_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False) |
There was a problem hiding this comment.
I suspect this happens a few times 😅 not to worry but maybe we need another pass to check
There was a problem hiding this comment.
Well the tests only complained on this one, so all models with a _keep_in_fp32_modules(_strict) should be safe at least - for others, I guess/hope we have tests about the dtype in general
| # Make sure the modules correctly exist if the flag is active | ||
| if model._keep_in_fp32_modules is None and model._keep_in_fp32_modules_strict is None: | ||
| self.skipTest( | ||
| reason=f"{model_class.__name__} has no _keep_in_fp32_modules nor _keep_in_fp32_modules_strict attribute defined" | ||
| ) |
There was a problem hiding this comment.
Just a thought, unsure: What if we had top level models where only a subpart of that model has these flags, e.g. a VLM where the text model has flags but otherwise not. Would these be also caught here
There was a problem hiding this comment.
That's a good point that should probably come in a following PR (see also my comment on other question). Currently, most vlms do
if language_model._keep_in_fp32_modules is not None:
self._keep_in_fp32_modules.extend(language_model._keep_in_fp32_modules)which is very not ideal but works...
| config, _ = self.model_tester.prepare_config_and_inputs_for_common() | ||
| for model_class in self.all_model_classes: | ||
| with self.subTest(model_class.__name__): | ||
| model = model_class(copy.deepcopy(config)) |
There was a problem hiding this comment.
Is there a reason we have to construct the model (previously it was inferred from the class)
There was a problem hiding this comment.
Because of this comment and behavior!
# Overwrite the class attribute to make it an instance attribute, so models like
# `InstructBlipForConditionalGeneration` can dynamically update it without modifying the class attribute
# when a different component (e.g. language_model) is used.
self._keep_in_fp32_modules = copy.copy(self.__class__._keep_in_fp32_modules)
self._keep_in_fp32_modules_strict = copy.copy(self.__class__._keep_in_fp32_modules_strict)
Cyrilvallez
added
the
for patch
4 participants
What does this PR do?
As per the title. #41580 broke the
keep_in_fp32_modulesflag as it's supposed to be used only with fp16, not bf16.I added very strict tests on this to avoid name clashing/typos etc!