Generic classifier builder #2356

SalmanMohammadi · 2025-02-07T00:57:25Z

This PR adds support for a generic classifier builder. Currently we require explicit classifier builders to be defined for every model which we wish to use as a classifier. This is costly to maintain.

We can instead define a generic classifier builder which simply adapts the final layer of any decoder-based model (how strict should this contract be?) to the desired number of classes.

This PR does not include support for PEFT-based classifier model builders. This is because there are no use-cases in the library for such model builders at this time.

pytorch-bot · 2025-02-07T00:57:29Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2356

📄 Preview Python docs built from this PR

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

⏳ No Failures, 6 Pending

As of commit c96a871 with merge base 32d195c ():
💚 Looks good so far! There are no failures yet. 💚

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

RdoubleA · 2025-02-07T02:35:32Z

torchtune/models/common.py

+    Returns:
+        nn.Module: The classifier model.
+    """
+    model = _get_component_from_path(base_model)(**base_model_kwargs)


if this is bound by the models in our library, then we would need to test that we can successfully create a classifier model for every model in our library, and any future models that are added. Can users pass in a custom model path? If that's not the intention, then we'll need to add some more gating. For example, if we check for if it's a TransformerDecoder instance that would easily keep this constrained to models that we know with confidence can be made into classifiers.

Great points, @RdoubleA.

Can users pass in a custom model path?

I think we only ever claim to support models in our library, right? So this would follow the same contract.

we would need to test that we can successfully create a classifier model for every model in our library, and any future models that are added

I'm curious why it wouldn't be sufficient to write an equally generic test against TransformerDecoder? I feel like the idea behind this builder is that we shouldn't have to define (or test) classifier builders for every model in our library - rather, I would think testing it against the generic cases we think could come up might be sufficient: model.output: TiedLinear, model.output: Linear, and maybe (model.output: TransformerDecoder.output). What do you think?

For example, if we check for if it's a TransformerDecoder instance that would easily keep this constrained to models that we know with confidence can be made into classifiers.

I think we agree here, if we can rely on any TransformerDecoder satisfying the requirements for this builder i.e. that it uses some linear output projection (which is a pretty flexible condition), then we should be able to use it as a classifier model. I think this would include vision models which use a TransformerDecoder as a decoder.

Yep agreed, if we test against TransformerDecoder that should satisfy all the conditions I mentioned

codecov-commenter · 2025-03-05T17:01:12Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 63.21%. Comparing base (67a8706) to head (ecf2e72).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2356      +/-   ##
==========================================
- Coverage   65.47%   63.21%   -2.27%     
==========================================
  Files         374      376       +2     
  Lines       22280    22316      +36     
==========================================
- Hits        14587    14106     -481     
- Misses       7693     8210     +517

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

felipemello1 · 2025-03-17T15:30:02Z

torchtune/models/common.py

+    if isinstance(model.output, nn.Linear):
+        del model.output.weight
+        if hasattr(model.output, "bias"):
+            del model.output.bias
+    model.output = nn.Linear(model.head_dim * model.num_heads, num_classes, bias=False)
+    return model


this wont work for TiedEmbedding layer.

Should we use

if hasattr(model, 'output')

Instead of:

if isinstance(model.output, nn.Linear)

also, for multimodal, it would have to be model.decoder.output, instead of model.output

Should we use
if hasattr(model, 'output')

Hmm so the logic would be something like -

if hasattr(model, 'output'): del model.output.weight if hasattr(model.output, "bias"): del model.output.bias

In the TiedEmbedding case, we will have model.output: TiedLinear, but model.output will have no weight/bias parameters to delete, so we won't need to enter this block right? Let me know if I'm missing something or misunderstanding here : )

also, for multimodal, it would have to be model.decoder.output, instead of model.output

I agree, but similar to PEFT, we don't support any use-cases for MM classifier models at the moment, so I thought I'd keep the design super simple and directed. What do you think?

sorry, still not clear to me why not just replace the output head if there is one, instead of deleting the weights first.

Just do:

if hasattr(model, 'output'): model.output = new_head()

I agree, but similar to PEFT, we don't support any use-cases for MM classifier models at the moment, so I thought I'd keep the design super simple and directed. What do you think?

I dont think it hurts to add and extra "else" to cover this case, if someone tries to experiments with it. If you are strongly against it, maybe we could raise an warning?

if hasattr(model, 'output'): do_something: else: raise "Could not find model.output"

felipemello1 · 2025-03-17T15:32:15Z

recipes/configs/mistral/7B_full_ppo_low_memory.yaml

+  _component_: torchtune.models.common.classifier_model
+  base_model_path: torchtune.models.mistral.mistral_7b


this works, but if we have nested instantiation, maybe a better design would be:

reward_and_value_model: _component_: torchtune.models.common.classifier_model model: _component_: torchtune.models.mistral.mistral_7b args

Do you mean changing base_model_path to model here? This sounds good to me : )

edited so that it renders a bit better

pbontrager · 2025-03-26T15:19:55Z

Instead of adding a common.py function here, could this be a classifier.py file in torchtune/modules?

joecummings · 2025-03-26T16:41:49Z

recipes/configs/mistral/7B_full_ppo_low_memory.yaml

@@ -46,17 +46,18 @@ policy_model:
 # we need to manually build the mistral classifier model
 # because our reward model checkpoint has a larger vocabulary size (due to an added padding token)
 reward_and_value_model:
-  _component_: torchtune.models.mistral._component_builders.mistral_classifier
+  _component_: torchtune.modules.classifier.classifier_model


felipemello1 · 2025-03-26T18:26:54Z

torchtune/modules/classifier.py

+    if hasattr(model, "output"):
+        del model.output
+        model.output = nn.Linear(
+            model.head_dim * model.num_heads, num_classes, bias=False
+        )
+    elif hasattr(model, "decoder") and hasattr(model.decoder, "output"):
+        del model.decoder.output
+        model.decoder.output = nn.Linear(
+            model.decoder.head_dim * model.decoder.num_heads, num_classes, bias=False
+        )
+    else:
+        raise ValueError("Could not find a valid output layer to adapt.")
+    return model


looks good! Not for this PR, but this wont work with PEFT. But i dont think we want to go there. Maybe we could add a warning: "found LoRA module, this is not supported, replacing it with a new head", but this might be an overkill. Any thoughts?

I can update the docstring but I think adding checks to cover PEFT cases might be overkill : )

felipemello1

will approve after CLI is green. Thanks for this PR!!

joecummings

This looks good! Can we file a TODO to add a small tutorial or something?

joecummings · 2025-03-27T14:47:03Z

torchtune/modules/classifier.py

+) -> Union[TransformerDecoder, nn.Module]:
+    """
+    Create a classifier model from a base model by adapting the output layer.
+    This builder does not support models which apply PEFT to the output layer.


nit: make this a Note:

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 7, 2025

RdoubleA reviewed Feb 7, 2025

View reviewed changes

SalmanMohammadi mentioned this pull request Mar 6, 2025

Enable PPO on Intel XPU using a tiny model #2446

Open

3 tasks

SalmanMohammadi requested a review from RdoubleA March 15, 2025 12:13

SalmanMohammadi marked this pull request as ready for review March 15, 2025 19:07

adding generic classifier builder

48786ea

SalmanMohammadi force-pushed the classifier_builder branch from 452d4b4 to 48786ea Compare March 15, 2025 19:11

SalmanMohammadi added 3 commits March 15, 2025 19:16

updating docs

6e65ac2

Merge branch 'main' into classifier_builder

4731204

updating docs

2b27089

felipemello1 reviewed Mar 17, 2025

View reviewed changes

SalmanMohammadi added 3 commits March 19, 2025 09:33

adding MM support

382d96e

linting

ab45c6f

fixing MM case

94a38f8

moving to modules

b92e2ae

joecummings reviewed Mar 26, 2025

View reviewed changes

SalmanMohammadi force-pushed the classifier_builder branch from 4df26a2 to 238a3ca Compare March 26, 2025 17:04

SalmanMohammadi added 2 commits March 26, 2025 17:07

fixing docs/imports

1ad8438

import shenanigans

244dd07

SalmanMohammadi force-pushed the classifier_builder branch from 238a3ca to 244dd07 Compare March 26, 2025 17:09

the last docstring fix - i swear

1a16642

SalmanMohammadi requested review from joecummings and felipemello1 March 26, 2025 17:11

Merge branch 'main' into classifier_builder

9034301

felipemello1 reviewed Mar 26, 2025

View reviewed changes

updating docstring

06d52eb

SalmanMohammadi mentioned this pull request Mar 26, 2025

Classifier builder v2 #2527

Closed

SalmanMohammadi added 8 commits March 27, 2025 11:29

testing CI

7638341

testing CI

80462ad

removing import

cd51736

removing all imports

a6e20a9

reordering tests

ecafd1a

reordering tests

7fc1abb

local seed

3fc436a

local seed

de1e66d

joecummings approved these changes Mar 27, 2025

View reviewed changes

updating docs

c96a871

SalmanMohammadi merged commit 0afea1a into pytorch:main Mar 27, 2025
17 checks passed

SalmanMohammadi deleted the classifier_builder branch March 27, 2025 16:07

This was referenced Mar 27, 2025

Deprecating classifier builders #2532

Merged

Classifiers (reward models) in torchtune #2320

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Generic classifier builder #2356

Generic classifier builder #2356

SalmanMohammadi commented Feb 7, 2025 •

edited

Loading

pytorch-bot bot commented Feb 7, 2025 •

edited

Loading

RdoubleA Feb 7, 2025

SalmanMohammadi Feb 7, 2025

RdoubleA Feb 7, 2025

codecov-commenter commented Mar 5, 2025

felipemello1 Mar 17, 2025 •

edited

Loading

SalmanMohammadi Mar 17, 2025 •

edited

Loading

felipemello1 Mar 17, 2025

felipemello1 Mar 17, 2025 •

edited

Loading

SalmanMohammadi Mar 17, 2025

felipemello1 Mar 17, 2025

pbontrager commented Mar 26, 2025

joecummings Mar 26, 2025

felipemello1 Mar 26, 2025 •

edited

Loading

SalmanMohammadi Mar 26, 2025

felipemello1 left a comment

joecummings left a comment

joecummings Mar 27, 2025

		_component_: torchtune.models.common.classifier_model
		base_model_path: torchtune.models.mistral.mistral_7b

Generic classifier builder #2356

Generic classifier builder #2356

Conversation

SalmanMohammadi commented Feb 7, 2025 • edited Loading

pytorch-bot bot commented Feb 7, 2025 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2356

⏳ No Failures, 6 Pending

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov-commenter commented Mar 5, 2025

Codecov Report

felipemello1 Mar 17, 2025 • edited Loading

Choose a reason for hiding this comment

SalmanMohammadi Mar 17, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 Mar 17, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbontrager commented Mar 26, 2025

Choose a reason for hiding this comment

felipemello1 Mar 26, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipemello1 left a comment

Choose a reason for hiding this comment

joecummings left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SalmanMohammadi commented Feb 7, 2025 •

edited

Loading

pytorch-bot bot commented Feb 7, 2025 •

edited

Loading

felipemello1 Mar 17, 2025 •

edited

Loading

SalmanMohammadi Mar 17, 2025 •

edited

Loading

felipemello1 Mar 17, 2025 •

edited

Loading

felipemello1 Mar 26, 2025 •

edited

Loading