Skip to content
This repository was archived by the owner on Sep 10, 2025. It is now read-only.

Push torchtune install deeper in generate._gen_model_input and undo import order change for etmodel #1542

Merged
merged 1 commit into from
May 12, 2025
Merged

Push torchtune install deeper in generate._gen_model_input and undo import order change for etmodel #1542

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 8 additions & 7 deletions torchchat/generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -869,13 +869,6 @@ def _gen_model_input(
max_new_tokens: Optional[int] = None,
max_seq_len: Optional[int] = 2048,
) -> Tuple[torch.Tensor, Optional[Dict[str, Any]]]:
# torchtune model definition dependencies
from torchtune.data import Message, padded_collate_tiled_images_and_mask
from torchtune.models.llama3_2_vision._model_builders import (
llama3_2_vision_transform,
)
from torchtune.training import set_default_dtype

"""
Convert prompt and image prompts into consumable model input args.

Expand Down Expand Up @@ -911,6 +904,14 @@ def _gen_model_input(
return encoded, None

# Llama 3.2 11B
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we sure that everything below is for llama 3.2 11b? If so good to go, otherwise we need an if-else.

Will we in the future have more possibilities here in this func? If so, we list them with an if-else or switch-case chain.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yup everything below is 11B.
If there's more models that require casing this function as a whole should get refactored at that time


# torchtune model definition dependencies
from torchtune.data import Message, padded_collate_tiled_images_and_mask
from torchtune.models.llama3_2_vision._model_builders import (
llama3_2_vision_transform,
)
from torchtune.training import set_default_dtype

assert (
image_prompts is None or len(image_prompts) == 1
), "At most one image is supported at the moment"
Expand Down
4 changes: 2 additions & 2 deletions torchchat/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1054,13 +1054,13 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

try:
# For llama::sdpa_with_kv_cache.out, preprocess ops
from executorch.extension.llm.custom_ops import custom_ops # no-qa
from executorch.extension.pybindings import portable_lib as exec_lib

# ET changed the way it's loading the custom ops so it's not included in portable_lib but has to be loaded separately.
# For quantized_decomposed ops
from executorch.kernels import quantized # no-qa
# For llama::sdpa_with_kv_cache.out, preprocess ops
from executorch.extension.llm.custom_ops import custom_ops # no-qa
Comment on lines +1062 to +1063
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: what's the difference this change made?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was the original order prior to the other PR
We preserve the order due to how extension libraries interact (we want custom_ops to overwrite potential conflicts)


class PTEModel(nn.Module):
def __init__(self, config, path) -> None:
Expand Down
Loading