CLIP Text Encoder #1969

calvinpelletier · 2024-11-07T23:23:35Z

Context

What is the purpose of this PR? Is it to

add a new feature
fix a bug
update tests and/or documentation
other (please add here)

Please link to any issues this PR addresses.

Changelog

What are the changes made in this PR?

Implement the CLIP tokenizer
Implement the CLIP text encoder

Test plan

Minimal code to run the CLIP text encoder e2e:
(first download CLIP weights: tune download openai/clip-vit-large-patch14 --output-dir /tmp/clip-vit-large-patch14 --ignore-patterns None)

import torch

from torchtune.models.clip._model_builders import clip_tokenizer, clip_text_encoder_large
from torchtune.training.checkpointing._checkpointer import FullModelHFCheckpointer

# build tokenizer and text encoder
tokenizer = clip_tokenizer()
encoder = clip_text_encoder_large()

# load weights
encoder.load_state_dict(FullModelHFCheckpointer(
    "/tmp/clip-vit-large-patch14",
    ["model.safetensors"],
    "CLIP_TEXT",
    "/tmp/torchtune-clip-vit-large-patch14",
).load_checkpoint()["model"])
encoder = encoder.to(torch.bfloat16).cuda().eval()

# run
text = [
    "a cow jumping over the moon",
    "a helpful AI assistant",
]
tokens = tokenizer(text)
encoding = encoder(tokens.cuda())

Checked parity with the HF CLIP tokenizer and text encoder as implemented here: MSE between the encoder outputs for on a batch of 32 test strings = 3.55e-5

Tokenization speed for 32 test strings

OpenAI: 0.0342
HuggingFace: 0.0416
TorchTune: 0.0195

Encoding speed for a single batch of 32 test strings:

HuggingFace: 1.699
TorchTune: 0.0173

Encoding speed for 1000 batches of 32 test strings:

HuggingFace: 30.498
TorchTune: 11.713

Checklist

run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
add unit tests for any new functionality
- TODO: unit test for encoder
update docstrings for any new or updated methods or classes
run unit tests via pytest tests
run recipe tests via pytest tests -m integration_test
manually run any new or modified recipes with sufficient proof of correctness
include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

I did not change any public API
I have added an example to docs or docstrings

pytorch-bot · 2024-11-07T23:23:39Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit 5aa7c9f with merge base 9bafd16 ():
💚 Looks good so far! There are no failures yet. 💚

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

tests/torchtune/models/clip/test_clip_tokenizer.py

torchtune/models/clip/_model_builders.py

torchtune/models/clip/_text_encoder.py

calvinpelletier · 2024-11-07T23:39:54Z

torchtune/models/clip/_tokenizer.py

+        self.sot_token = self.encoder["<|startoftext|>"]
+        self.eot_token = self.encoder["<|endoftext|>"]
+        self.pad_token = self.eot_token


should these be configurable in the constructor like we do for other tokenizers? I don't really think it's needed imo

They might not need to be, but I think it's more readable when all of the special tokens are listed at the top. I'm going to ask @RdoubleA to review this file specifically as he has the most experience with how tokenizers are used in torchtune.

I think keeping these here is fine since this is operating as a base tokenizer and not a model tokenizer with tokenize_messages

torchtune/models/clip/_tokenizer.py

codecov-commenter · 2024-11-07T23:48:44Z

Codecov Report

Attention: Patch coverage is 92.20183% with 17 lines in your changes missing coverage. Please review.

Project coverage is 67.82%. Comparing base (4b6877a) to head (d5b7f98).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
torchtune/models/clip/_convert_weights.py	36.36%	7 Missing ⚠️
torchtune/models/clip/_tokenizer.py	93.39%	7 Missing ⚠️
torchtune/models/clip/_model_builders.py	88.88%	1 Missing ⚠️
torchtune/models/clip/_text_encoder.py	95.83%	1 Missing ⚠️
torchtune/training/checkpointing/_checkpointer.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1969      +/-   ##
==========================================
+ Coverage   67.51%   67.82%   +0.30%     
==========================================
  Files         318      323       +5     
  Lines       17684    17897     +213     
==========================================
+ Hits        11940    12139     +199     
- Misses       5744     5758      +14

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

torchtune/models/clip/_tokenizer.py

pbontrager

This is a great PR! I left some comments around some standard patterns we try to follow but aside from that, this looks very solid and clean.

torchtune/models/clip/_tokenizer.py

pbontrager · 2024-11-11T16:45:03Z

torchtune/modules/activations.py

+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.sigmoid(1.702 * x)


How does this affect parity values and speed (as opposed to classic GeLU)? I'd prefer not to add this extra module if it's only ever going to be used by CLIP.

the speedup of quickgelu is negligible but it's needed because using default gelu results in significantly worse parity (output MSE: 0.16 vs 0.00003)

Let's move this to the clip components builder file for now then. Just a function "def quick_gelu"

moved it to the clip components file but left it as a module because the activation arg of clip_mlp/FeedForward expects a module (I could move it to the clip/_text_encoder.py file instead if that's better)

pbontrager · 2024-11-11T16:47:47Z

torchtune/training/checkpointing/_checkpointer.py

@@ -488,6 +489,10 @@ def load_checkpoint(self) -> Dict[str, Any]:
                    "supported_aspect_ratios", None
                ),
            )
+        elif self._model_type == ModelType.CLIP_TEXT:


I'd probably prefer the model type to just be CLIP and the convert function to be clip_hf_to_tune. We don't support the vision model yet but we plan to and I think it'd be simpler to have just one CLIP version. The clip_hf_to_tune model can raise and error for now if someone attempts to load in the vision model.

how would this work? the clip_hf_to_tune function needs to know whether it's creating torchtune parameters for the text or the vision model. and the only relevant information that the checkpointer has access to is the model_type, right? so the model type needs to specify text or vision?

We'd just have a single convert mapping dictionary and map what happens to be there. I don't think the function actually needs to know which fraction of the weights it's converting.

Every param will be there since HF combines the vision and text params. So the function needs to know what to ignore

torchtune/utils/_download.py

torchtune/models/clip/_tokenizer.py

torchtune/models/clip/_model_builders.py

tests/torchtune/models/clip/test_clip_tokenizer.py

pbontrager

I'd still like the model type to be unified to just CLIP to keep the list shorter. But I'll approve when @RdoubleA signs off on the tokenizer.

torchtune/models/clip/_model_builders.py

torchtune/models/clip/_text_encoder.py

pbontrager · 2024-11-14T15:58:47Z

torchtune/modules/activations.py

+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.sigmoid(1.702 * x)


Let's move this to the clip components builder file for now then. Just a function "def quick_gelu"

pbontrager · 2024-11-14T16:00:18Z

torchtune/training/checkpointing/_checkpointer.py

@@ -488,6 +489,10 @@ def load_checkpoint(self) -> Dict[str, Any]:
                    "supported_aspect_ratios", None
                ),
            )
+        elif self._model_type == ModelType.CLIP_TEXT:


We'd just have a single convert mapping dictionary and map what happens to be there. I don't think the function actually needs to know which fraction of the weights it's converting.

RdoubleA · 2024-11-14T20:33:29Z

torchtune/models/clip/_tokenizer.py

+        truncate (bool): whether to truncate the text when longer than max_seq_len
+    """
+
+    def __init__(self, path: PathLike, max_seq_len: int = 77, truncate: bool = True):


nit: does PathLike include both str and pathlib? leaning towards keeping it just str which is consistent with our other tokenizers

RdoubleA · 2024-11-14T20:34:34Z

torchtune/models/clip/_tokenizer.py

+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+
+        merges = []
+        with open(path, encoding="utf-8") as f:


nit: let's move this logic to a method

RdoubleA · 2024-11-14T20:59:42Z

torchtune/models/clip/_tokenizer.py

+            .replace("</w>", " ")
+        )
+
+    def __call__(self, texts: List[str]) -> torch.Tensor:


we should return token ids as List[int] instead of tensor, the collater usually converts to tensor. also, usually the __call__ method would be for a single sample (see example: https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3/_tokenizer.py#L330) instead of a batch of texts. The tokenizer will be in the dataset, and will automatically call it on each sample individually.

In this case, maybe the call method would only call encode, or you don't need the call method at all. It depends on if you'll be using the same SFTDataset or a different text + image dataset abstraction that won't rely on the Message structure

RdoubleA · 2024-11-14T21:03:39Z

torchtune/models/clip/_tokenizer.py

+        if token in self.cache:
+            return self.cache[token]
+
+        word = tuple(token[:-1]) + (token[-1] + "</w>",)


maybe </w> should be a class attribute or constant since you are using it everywhere

RdoubleA · 2024-11-14T21:05:57Z

torchtune/models/clip/_tokenizer.py

+        return result
+
+    def _bpe(self, token):
+        if token in self.cache:


will keeping a cache lead to exploding memory usage as we tokenize more and more in the same script?

an alternative could be to use the decorator @lru_cache, but not sure if that will degrade performance too much

RdoubleA · 2024-11-14T21:07:21Z

torchtune/models/clip/_tokenizer.py

+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):


I know this is straight from the official implementation, but could you add some high level comments to explain what's happening in these loops?

RdoubleA · 2024-11-14T21:08:02Z

torchtune/models/clip/_tokenizer.py

+    return dict(zip(bs, cs))
+
+
+def _get_pairs(word):


type annotations

RdoubleA · 2024-11-14T21:08:09Z

torchtune/models/clip/_tokenizer.py

+        return word
+
+
+def _bytes_to_unicode():


return type annotations

RdoubleA · 2024-11-14T21:08:17Z

torchtune/models/clip/_tokenizer.py

+            result[i, : len(tokens)] = torch.tensor(tokens)
+        return result
+
+    def _bpe(self, token):


type annotations

RdoubleA · 2024-11-14T21:08:57Z

torchtune/models/clip/_tokenizer.py

+        self.sot_token = self.encoder["<|startoftext|>"]
+        self.eot_token = self.encoder["<|endoftext|>"]
+        self.pad_token = self.eot_token


I think keeping these here is fine since this is operating as a base tokenizer and not a model tokenizer with tokenize_messages

calvinpelletier added 3 commits November 7, 2024 13:16

CLIP tokenizer and text encoder

f40d879

Merge remote-tracking branch 'origin/main' into clip_text

5dcf0d0

formatting

0a070af

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 Nov 7, 2024

calvinpelletier commented Nov 7, 2024

View reviewed changes

calvinpelletier requested a review from pbontrager November 7, 2024 23:46

calvinpelletier commented Nov 8, 2024

View reviewed changes

torchtune/models/clip/_tokenizer.py Outdated Show resolved Hide resolved

pbontrager reviewed Nov 11, 2024

View reviewed changes

calvinpelletier added 6 commits November 11, 2024 13:54

switching to hf vocab file

8334463

remove dependency on ftfy

d43f5b0

clip text encoder unit test

a3c90ac

address comments

1dbe939

move __call__

e6b3d19

merge

d501903

pbontrager reviewed Nov 14, 2024

View reviewed changes

calvinpelletier added 3 commits November 14, 2024 11:17

addressing comments

c4e700b

Merge remote-tracking branch 'origin/main' into clip_text

d5b7f98

moving quickgelu

5aa7c9f

RdoubleA reviewed Nov 14, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

CLIP Text Encoder #1969

CLIP Text Encoder #1969

calvinpelletier commented Nov 7, 2024 •

edited

Loading

pytorch-bot bot commented Nov 7, 2024 •

edited

Loading

calvinpelletier Nov 7, 2024

pbontrager Nov 11, 2024

RdoubleA Nov 14, 2024

codecov-commenter commented Nov 7, 2024 •

edited

Loading

pbontrager left a comment

pbontrager Nov 11, 2024

calvinpelletier Nov 11, 2024

pbontrager Nov 14, 2024

calvinpelletier Nov 14, 2024

pbontrager Nov 11, 2024

calvinpelletier Nov 11, 2024

pbontrager Nov 14, 2024

calvinpelletier Nov 14, 2024

pbontrager left a comment

pbontrager Nov 14, 2024

pbontrager Nov 14, 2024

RdoubleA Nov 14, 2024 •

edited

Loading

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

RdoubleA Nov 14, 2024

CLIP Text Encoder #1969

Are you sure you want to change the base?

CLIP Text Encoder #1969

Conversation

calvinpelletier commented Nov 7, 2024 • edited Loading

Context

Changelog

Test plan

UX

pytorch-bot bot commented Nov 7, 2024 • edited Loading

🔗 Helpful Links

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

❗ 1 Active SEVs

✅ No Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov-commenter commented Nov 7, 2024 • edited Loading

Codecov Report

pbontrager left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbontrager left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RdoubleA Nov 14, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

calvinpelletier commented Nov 7, 2024 •

edited

Loading

pytorch-bot bot commented Nov 7, 2024 •

edited

Loading

codecov-commenter commented Nov 7, 2024 •

edited

Loading

RdoubleA Nov 14, 2024 •

edited

Loading