Add custom resize_token_embeddings method to OLMoForCausalLM (#491)

[djliden]

This PR addresses #491 by adding a custom resize_token_embeddings method to override the corresponding method from the PreTrainedModel base class. The revised method resizes self.config.embedding_size and self.model.config.embedding_size rather than resizing the vocab_size as done by the base class's method.

Here's the result with the new method; this resulted in an error previously (see #491).

!pip install -e .
!pip install accelerate
dbutils.library.restartPython()

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import hf_olmo

model_ckpt = "allenai/OLMo-1B"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = AutoModelForCausalLM.from_pretrained(
    model_ckpt,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

print("Before resizing:")
print(f"Vocabulary size: {model.config.vocab_size}")
print(f"Embedding size: {model.config.embedding_size}")
print(f"Model embedding dimension: {model.get_input_embeddings().embedding_dim}")

special_tokens = ["<|im_start|>", "<|im_end|>"]
tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
model.resize_token_embeddings(len(tokenizer))

print("\nAfter resizing:")
print(f"Vocabulary size: {model.config.vocab_size}")
print(f"Embedding size: {model.config.embedding_size}")
print(f"Model embedding dimension: {model.get_input_embeddings().embedding_dim}")

input_text = "This is a dummy input text."
input_ids = tokenizer.encode(input_text, return_tensors="pt")
labels = input_ids.clone()

input_ids = input_ids.to(model.device)
labels = labels.to(model.device)

with torch.no_grad():
    outputs = model(input_ids, labels=labels)

print("\nModel run completed without errors.")

Resulting in:

Before resizing:
Vocabulary size: 50280
Embedding size: 50304
Model embedding dimension: 2048

After resizing:
Vocabulary size: 50280
Embedding size: 50282
Model embedding dimension: 2048

Model run completed without errors.

Note that this does not update the vocab_size. The intended purpose of having separate vocab_size and embedding_size is not entirely clear to me yet, but it seems like the intent is to allow separate vocab_size and embedding_size, as long as embedding_size > vocab_size. A few further options for addressing this issue more completely include:

• update the new method to check whether config.embedding_size is defined, and update vocab_size if it isn't. This corresponds to how the wte layer is defined in the first place (link)
• add a separate function for resizing the vocabulary
• treat vocab size vs. embedding size discrepancies as a separate issue and review the current change as-is

Please let me know if you have any questions or suggestions regarding this PR.

[djliden]

@2015aroras @AkshitaB I'm not able to assign reviewers, so flagging for review as requested in #491 . Thanks!

[2015aroras]

Vocab size captures how many different tokens a tokenizer can produce. Embedding size captures how many different integers the embedding layer can produce an embedding for. Theoretically speaking, the vocab and embedding sizes would be the same, since the embedding layer should have exactly 1 embedding per token and no extra embeddings.

In practice, not having an embedding size that is a multiple of 128 can can slight performance degradation (broadly speaking, having factors of 2 seems to be beneficial). Thus we use a bigger embedding size than needed, resulting in there being embeddings for some integers that are not valid tokens. Hopefully that makes sense.

The PR itself looks good! Just requesting some small changes.

[2015aroras]

This comment seems to be the exact same as the one from the parent class. If it is, then it's not needed here (code editors and IDEs should be able to show the parent class's comment when one hovers over resize_token_embeddings). Having the comment somewhat implies that we are defining what the method does, but it is the parent doing it.

[2015aroras]

Could you add a warning for when the new embedding size is less than the vocab size? Something like "Resizing token embeddings to size <size>, which is less than the vocab size <vocab size> of the tokenizer.

[djliden]

@2015aroras I added a warning if the new embedding size is < the vocab size. I could see an argument for making it an error for consistency with this, or omitting entirely since it will result in an error as soon as you try to do anything with it anyway. Though I think it also works as-is.

I left the comment alone as it does change vocab size to embedding size (so not the same as base class), but please let me know if you'd like it amended regardless. Thanks!

[2015aroras]

@2015aroras I added a warning if the new embedding size is < the vocab size. I could see an argument for making it an error for consistency with this, or omitting entirely since it will result in an error as soon as you try to do anything with it anyway. Though I think it also works as-is.

I leaned towards warning instead of error since the user is actively modifying the embedding size. I don't feel strongly about this and resizing embeddings is a relatively uncommon scenario, so any of the 3 options is ok.

I left the comment alone as it does change vocab size to embedding size (so not the same as base class), but please let me know if you'd like it amended regardless. Thanks!

I didn't notice the vocab size to embedding size change! I think it's ok to leave it as is, though you could describe how the method differs from PreTrainedModel.resize_token_embeddings.

[djliden]

@2015aroras I updated the docstring noting the difference, and I'm happy leaving it as a warning. Thanks!

[2015aroras]

Just small things, otherwise looking good!

[2015aroras]

We tend to use log = logging.getLogger(__name__) and log.warn(...) to log warnings now (e.g.

OLMo/scripts/train.py

Line 45 in 417af0e

log.warning(

).

[2015aroras]

nit: Not enough indentation on this line

[djliden]

@2015aroras Changed the warning format and fixed the indentation. Thanks!

[2015aroras]

LGTM! Just fix the code style using Ruff and then you can merge this.

[djliden]

@2015aroras done!

[djliden]

@2015aroras Thanks for the reviews/feedback/etc! I don't have write access to merge this PR myself—from my side, both this PR and #503 are ready to go. Please let me know if there's anything else needed, otherwise feel free to merge when you have a chance. Appreciate your help!