Repository navigation
[gemma3] fix bidirectional image mask - #39396
Conversation
|
run-slow: gemma3 |
|
This comment contains run-slow, running the specified jobs: models: ['models/gemma3'] |
|
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. |
|
run-slow: gemma3 |
|
This comment contains run-slow, running the specified jobs: models: ['models/gemma3'] |
|
Just realized the |
|
run-slow: gemma3 |
|
This comment contains run-slow, running the specified jobs: models: ['models/gemma3'] |
| image_group_ids_at_kv_idx = torch.where(kv_idx < image_group_ids.shape[1], image_group_ids_at_kv_idx, -1) | ||
|
|
||
| is_image_block = (token_type_ids[batch_idx, q_idx] == 1) & (token_type_ids_at_kv_idx == 1) | ||
| same_image_block = image_group_ids[batch_idx, q_idx] == image_group_ids_at_kv_idx |
There was a problem hiding this comment.
we could use only token_type_ids but since vmap doesn't support slicing, it always fails with an error. I tried out with different torch ops, none of them worked. Passing image_group_ids was the most straightforward woraround
|
Can be reviewed now, slow tests are failing and I'll fix them tomorrow. No idea how to get Edit: yeah, we need to login with our own access tokens after ssh-ing |
|
run-slow: gemma3 |
|
This comment contains run-slow, running the specified jobs: models: ['models/gemma3'] |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Can we have a fast test for this? 🤗
| image_group_ids_at_kv_idx = torch.where(kv_idx < image_group_ids.shape[1], image_group_ids_at_kv_idx, -1) | ||
|
|
||
| is_image_block = (token_type_ids[batch_idx, q_idx] == 1) & (token_type_ids_at_kv_idx == 1) | ||
| same_image_block = image_group_ids[batch_idx, q_idx] == image_group_ids_at_kv_idx |
|
Actually the slow tests should fail whenever it breaks, prob we didn't notice it because slow tests were failing for another reason already. I will try to make a fast test to check attentions |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: gemma3 |
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
* fix gemma3 mask * make compile happy, and use only torch ops * no full attention between images * update tests * fix tests * add a fast test
What does this PR do?
Fixes #39389. There were two issues. First, using
getattron dict will never look for a dict key thus we were skipping ontoken_type_ids_mask.Second,
kv_idxcan go beyondtoken_type_idslength because it includes the new tokens to be generated. Using a simple check onids.shapefails in vmap due to dynamic control flow. One option would be to create dummy padded token type ids usingkv_idxbut I think a simple try-except is also fine