Repository navigation
Add dots3-note Preview model support - #47844
miraclezqc wants to merge 26 commits into
Conversation
7af5a25 to
fe87862
Compare
|
hey @miraclezqc , thanks for the PR! Most of the team currently is off on summer vacations so the reviews are taking longer. We'll review it soon :) |
b70c2d9 to
d8e147f
Compare
@zucchini-nlp |
84acd71 to
9daa866
Compare
|
Hey! Happy to review, can you update the branch to fix the latest |
Thanks! The branch has been updated for the latest mlinter rules. It seems that the new CI run is blocked by the security gate. Could you please approve it? |
|
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. |
CI recapDashboard: View test results in Grafana |
|
Thanks for approving the workflow! All regular CI checks are now green, and the branch is ready for review. |
|
Hi @ArthurZucker @Rocketknight1 , the branch is now up to date with main, and all relevant local tests pass. We’d appreciate your review and an early merge if everything looks good. Thank you! |
|
Sorry for the delays @miraclezqc but I noticed that we are not using modular as intended at all, please check out https://huggingface.co/docs/transformers/modular_transformers The goal is to reuse existing modules as much as we can, this includes the rms norm, indexer, etc |
|
@vasqu Thanks! We’ve refactored the model to reuse existing DeepSeek, GLM DSA, and Qwen2-VL components, including RMSNorm and the indexer. Only Dots-specific logic remains custom, and all relevant tests pass. Could you please take another look and let us know if it is ready to merge? |
vasqu
left a comment
There was a problem hiding this comment.
I tried to go a bit more into details but this is still not super modular friendly: There are a lot of custom implementations where we could just reuse ours - e.g. the dsa is handled completely custom, the attn backend is also completely custom which would disallow any usage via vllm, etc.
What I want to tell with this: Custom implementations are very much still implemented in transformers oftentimes but they need to be aligned - imo the biggest example is RoPE in this case
| > [!NOTE] | ||
| > Loading encoded audio or video sources requires the optional `torchcodec` dependency. Native video preprocessing | ||
| > follows the training-time sampling pipeline and expands each video into timestamped image blocks with interleaved | ||
| > audio blocks. The `<|video_pad|>` marker is only an external prompt placeholder and is removed before tokenization; | ||
| > it is never passed to the model. Decoded frame arrays are supported as visual-only inputs because their original | ||
| > audio track is absent. |
There was a problem hiding this comment.
I think this is fairly common across all of our models so I think we can delete it but no strong opinion
| return model | ||
|
|
||
|
|
||
| def _resolve_weight_block_size(hf_quantizer, value: torch.Tensor) -> tuple[int, int]: |
There was a problem hiding this comment.
this looks unnecessary no? Can we revert to keep it inline and not a separate fn?
There was a problem hiding this comment.
cc @IlyasMoutawwakil when you have time to check this over
can you share the motivation here? Ig there is some (new) fp4 handling we need?
There was a problem hiding this comment.
still need an answer here, seems like something happens alongside the sharding so we don't have enough to properly unpack?
we also want to move to a more uniform api in #48058 so would help which fp4 format we have here
| pass | ||
|
|
||
|
|
||
| class Dots3NoteVisionRotaryEmbedding(nn.Module): |
| return attention_output, attention_weights | ||
|
|
||
|
|
||
| class Dots3NoteVisionAttention(nn.Module): |
There was a problem hiding this comment.
same here a lot of custom stuff we should not use and try to align with existing models
| # Unified multimodal model | ||
| # ----------------------------------------------------------------------------- | ||
| @auto_docstring | ||
| class Dots3NoteForCausalLM(Dots3NoteTextForCausalLM): |
There was a problem hiding this comment.
ForConditional - pls check the pattern in qwen vl
There was a problem hiding this comment.
or gemma4 and the like (omni)
|
|
||
|
|
||
| @auto_docstring | ||
| class Dots3NoteForConditionalGeneration(Dots3NoteForCausalLM): |
There was a problem hiding this comment.
Doesnt make much sense to make for causal lm and for conditional
Is the goal to have a text only variation, then pls check out qwen3_5 which does that
There was a problem hiding this comment.
Modular can also be applied for all processor related things
Didnt go through these but we have a lot of custom things we shouldnt need
Thanks @vasqu for the detailed review and all the helpful suggestions. We’ve gone through the comments and refactored the text, audio, vision, and preprocessing code to reuse existing Transformers components, removing redundant implementations along the way. We believe most suggestions are addressed. The remaining differences and follow-ups are:
We may still have missed a comment or a simpler reuse opportunity. Could you take another look and let us know what else needs to change before merging? Thanks again for your time and guidance. |
|
Checking tomorrow! 🤗 |
vasqu
left a comment
There was a problem hiding this comment.
Sorry for the delays, another round 🫡
| language-model head and generation interface. `Dots3NoteTextForCausalLM` provides the text-only variant; | ||
| `Dots3NoteForCausalLM` remains a compatibility name for the original multimodal checkpoints. |
There was a problem hiding this comment.
imo its fine to have the for causal lm only variation, we dont really use textforcausallm
There was a problem hiding this comment.
still need an answer here, seems like something happens alongside the sharding so we don't have enough to properly unpack?
we also want to move to a more uniform api in #48058 so would help which fp4 format we have here
| MISSING_VIDEO_PROCESSOR_MAPPING_NAMES = OrderedDict( | ||
| [ | ||
| ("cosmos3_omni", "Qwen3VLVideoProcessor"), | ||
| ("dots3_note", "Dots3NoteVideoProcessor"), |
There was a problem hiding this comment.
this is weird, should usually be detected into auto mappings no?
| ("dinat", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}), | ||
| ("dinov2", {"torchvision": "BitImageProcessor", "pil": "BitImageProcessorPil"}), | ||
| ("donut-swin", {"torchvision": "DonutImageProcessor", "pil": "DonutImageProcessorPil"}), | ||
| ("dots3_note", {"pil": "Dots3NoteImageProcessor"}), |
There was a problem hiding this comment.
same here should be detected in auto mappings
| pixel_values: torch.Tensor, | ||
| image_grid_thw: torch.Tensor, | ||
| **kwargs, | ||
| ) -> torch.Tensor: |
There was a problem hiding this comment.
it should return a model output class, check qwen vl
| return output.audio_embeds, output.audio_token_lengths | ||
|
|
||
| @staticmethod | ||
| def _merge_multimodal_embeddings( |
There was a problem hiding this comment.
lets use the existing placeholder mask pattern please
| chunk_sample_lengths (`torch.Tensor`, *optional*): Waveform sample count for each audio feature chunk. | ||
| chunk_token_lengths (`torch.Tensor`, *optional*): Encoder token count for each audio feature chunk. | ||
| audio_chunk_counts (`torch.Tensor`, *optional*): Number of feature chunks for each audio input. | ||
| audio_token_lengths (`torch.Tensor`, *optional*): Expected encoded token count for each audio input. | ||
| chunk_audio_indices (`torch.Tensor`, *optional*): Audio ownership metadata emitted by the processor. |
There was a problem hiding this comment.
yea im still unsure here because this feels overkill and might be better described with metadata like mm token ids and cu seqlens etc?
| if isinstance(image, Image.Image) and image.mode == "RGBA": | ||
| from ..idefics2 import image_processing_pil_idefics2 | ||
|
|
||
| return image_processing_pil_idefics2.convert_to_rgb(image) |
There was a problem hiding this comment.
I'd rather we copy paste because i dont think modular would properly work here to unfold the implementation
| if sample_end > sample_start: | ||
| waveform = np.ascontiguousarray(pcm[sample_start:sample_end].astype(np.float32) / 32768.0) | ||
| output.append({"type": "audio", "audio": waveform}) | ||
| return output |
There was a problem hiding this comment.
this looks way too custom, let's try to move it into our processor framework
E.g. video processor have sample frames functions we can override; loading any video should not be manually done as we load the video with metadata etc
|
Thanks @vasqu for another detailed review! We’ve pushed a further refactor. One detail worth clarifying:
Could you take another look? Thanks again for your patience and guidance! |
vasqu
left a comment
There was a problem hiding this comment.
Next round but this looks very good already! I have smaller concerns but the overall core is solid
| pyramid_num_routed (`list[int]` or `tuple[int, ...]`, *optional*): | ||
| Number of routed experts in each vision transformer layer. A non-positive value selects a dense MLP. | ||
| capacity_factor (`int`, *optional*, defaults to 2): | ||
| Number of experts selected per token, capped by the layer's expert count. |
There was a problem hiding this comment.
I have to check the vision model in detail but why not use the same moe conventions as we have it in the text, e.g. num_experts_per_tok (capacity_factor) and mlp_layer_types (pyramid_num_routed)
|
|
||
| model_type = "dots3_note_vision_encoder" | ||
| base_config_key = "vision_config" | ||
| attribute_map = {"num_heads": "num_attention_heads", "use_bias": "attention_bias"} |
There was a problem hiding this comment.
why do we need this attribute mapping? You can also let them live via kwargs being passed in post init (but we do not recognize it here)
| if kwargs.pop("is_causal", False): | ||
| raise ValueError("Dots 3 Note Preview vision attention requires is_causal=False") | ||
| if not kwargs.pop("pre_pixel_shuffle", True) or kwargs.pop("adapter_type", "patch_merger") != "patch_merger": | ||
| raise ValueError("Dots 3 Note Preview requires pre_pixel_shuffle=True and adapter_type='patch_merger'") |
There was a problem hiding this comment.
imo this is too overly cautious, let's just make sure our defaults are really fitting
if you have something that checks values, use validate_architecture
| attribute_map = { | ||
| "d_model": "hidden_size", | ||
| "encoder_attention_heads": "num_attention_heads", | ||
| "encoder_layers": "num_hidden_layers", | ||
| "encoder_ffn_dim": "intermediate_size", | ||
| } |
There was a problem hiding this comment.
same here if we can avoid using attribute mapping it would be nice
| self.head_dim = self.hidden_size // self.num_attention_heads | ||
| self.num_key_value_heads = self.num_attention_heads | ||
| self.hidden_act = "silu" | ||
| self.attention_bias = True |
There was a problem hiding this comment.
a bit weird, can be just set correctly as attribute no? (maybe happens in vision/text as well?)
| kwargs.update( | ||
| pixel_values=pixel_values, | ||
| pixel_values_videos=pixel_values_videos, | ||
| image_grid_thw=image_grid_thw, | ||
| video_grid_thw=video_grid_thw, | ||
| input_features=input_features, | ||
| chunk_sample_lengths=chunk_sample_lengths, | ||
| ) |
There was a problem hiding this comment.
let's rather inherit from another omni model and avoid this
There was a problem hiding this comment.
Rebump: no kwargs manipulation like this - rather write the complete forward again than this
| return_attention_mask=return_attention_mask, | ||
| **kwargs, | ||
| ) | ||
| self.dither = 0.0 |
There was a problem hiding this comment.
hmm this is a bit weird no? why is it not just an arg as the others?
| # Correct legacy Qwen2-VL defaults without overriding custom pixel limits. | ||
| if image_processor is not None: | ||
| if dict(image_processor.size) == _QWEN2_VL_IMAGE_DEFAULT_SIZE: | ||
| image_processor.size = SizeDict(**dict(_RELEASE_VISION_SIZE)) | ||
| if image_processor.temporal_patch_size == 2: | ||
| image_processor.temporal_patch_size = 1 | ||
| if video_processor is not None: | ||
| if dict(video_processor.size) == _QWEN2_VL_VIDEO_DEFAULT_SIZE: | ||
| video_processor.size = SizeDict(**dict(_RELEASE_VISION_SIZE)) | ||
| if video_processor.temporal_patch_size == 2: | ||
| video_processor.temporal_patch_size = 1 |
There was a problem hiding this comment.
imo should be updated on the hub instead no? we shouldnt need that workaround
|
|
||
| def convert_to_rgb(self, image): | ||
| return convert_to_rgb(image) | ||
|
|
|
Thanks @vasqu! We’ve gone through your comments and pushed another round of changes. A few notes:
Could you take another look when you have a chance? We’d appreciate your feedback on anything we may have missed. |
vasqu
left a comment
There was a problem hiding this comment.
Sorry that I'm doing so many comments again but we are definitely getting closer!
There are a lot of smaller details but overall the biggest thing is the tests now imo which should follow other similar models like gemma4 (so omni models)
| a shared vision encoder for images and videos and a Whisper-style audio encoder. Both encoders project their outputs | ||
| into the language model's hidden space before autoregressive text generation. | ||
|
|
||
| This integration supports BF16 checkpoints. |
There was a problem hiding this comment.
Mainly the indexer or completely not usable? I think we dequant the indexer but the rest could use fp8 no?
|
|
||
| def __post_init__(self, **kwargs): | ||
| if self.rope_parameters is None: | ||
| self.rope_parameters = {"rope_type": "axial", "rope_theta": 10_000.0} |
There was a problem hiding this comment.
Let's just use default_rope_type="axial" instead then we dont need this
| if pyramid_num_routed is None and self.mlp_layer_types is None: | ||
| pyramid_num_routed = [-1] * 25 + list(range(4, 65, 4)) + [64] | ||
| self.num_experts_per_tok = kwargs.pop("capacity_factor", self.num_experts_per_tok) | ||
| self.router_scaling_factor = kwargs.pop("router_scale", self.router_scaling_factor) |
There was a problem hiding this comment.
imo you can just do get if you want these values to surive serialization
| if (attention_backend := kwargs.pop("attention_backend", None)) is not None: | ||
| kwargs.setdefault("attn_implementation", attention_backend) |
There was a problem hiding this comment.
should not be done, we should rely on our set attn implementation etc and default to work out
| self.num_key_value_heads = self.num_attention_heads | ||
| self.rope_parameters = { | ||
| "rope_type": "default", | ||
| "partial_rotary_factor": 0.5, |
There was a problem hiding this comment.
Ig you only want the partial rotary factor, I think kwargs.setdefault("partial_rotary_factor", 0.5) should suffice
| # Initial support is inference-focused, so common coverage targets forward, cache, and generation. | ||
| self.is_training = False |
There was a problem hiding this comment.
imo can still be turned on, no resaon to turn it off. The indexer is the only part that has no proper aux loss but its known
| # Initial support is inference-focused, so common coverage targets forward, cache, and generation. | ||
| self.is_training = False | ||
|
|
||
| def get_config(self): |
There was a problem hiding this comment.
not needed if we have a proper init override
| parent=parent, | ||
| batch_size=2, | ||
| seq_length=7, | ||
| vocab_size=128, | ||
| hidden_size=32, | ||
| num_hidden_layers=2, | ||
| num_attention_heads=4, | ||
| num_key_value_heads=4, | ||
| intermediate_size=64, | ||
| max_position_embeddings=128, |
There was a problem hiding this comment.
a lot of common kwarg like these are already covered. pls check e.g. youtu on the style
| all_model_classes = (Dots3NoteForCausalLM,) if is_torch_available() else () | ||
| pipeline_model_mapping = {"text-generation": Dots3NoteForCausalLM} if is_torch_available() else {} | ||
| _is_stateful = True |
There was a problem hiding this comment.
| all_model_classes = (Dots3NoteForCausalLM,) if is_torch_available() else () | |
| pipeline_model_mapping = {"text-generation": Dots3NoteForCausalLM} if is_torch_available() else {} | |
| _is_stateful = True |
we should mark stateful in pretrained model if anything but I doubt it
|
|
||
|
|
||
| @require_torch | ||
| class Dots3NoteModelTest(unittest.TestCase): |
There was a problem hiding this comment.
Let's align with gemma4 style please
There are text only, vision+text, and audio+text tests similar to how you need it here + then finally a few integration tests
|
Thank you for your contribution 🤗! CI Security Gate — automatic approval blockedThis PR was not automatically approved for CI because the security gate failed. Possible reasons:
See the workflow run for the exact violations. A maintainer can review and manually approve CI if a finding is a false positive. |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, dots3_note |
|
No worries at all, and thanks for your patience! I’m still getting familiar with Transformers conventions, which has made the iterations slower. A few points remain:
I’m happy to go through another round and keep iterating until everything is resolved. |
Hi @vasqu, do you have any feedback on the latest changes? |
What does this PR do?
This PR is submitted by the official Dots-Studio team.
This PR adds native Transformers support for Dots 3 Note Preview, an inference-only mixture-of-experts multimodal causal language model supporting text, image, video, and audio inputs.
The implementation follows the modular Transformers workflow and includes the model configuration, processor, image processor, audio feature extractor, modeling implementation, auto-class registrations, documentation, conversion utilities, and model tests required for native integration.
The implementation includes:
128 x 128weight block size;This integration currently targets inference. Training-specific features such as gradient checkpointing are outside the scope of this PR.
No additional runtime dependency is introduced by this integration.
Validation
The following checks were run against the implementation in this PR:
tests/models/dots3_notetest directory;AutoModelForMultimodalLM;Both the BF16 and FP8 checkpoints passed the end-to-end text, image, video, and audio validation cases with finite outputs and successful greedy generation.
The corresponding public checkpoint links and finalized model organization metadata will be added after this PR is merged.
Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice.
As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues.
We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.
Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we
might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not
want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise
for reviewers and other users and will almost certainly get you blocked.
For more information, please read
CONTRIBUTING.md.Before submitting
Who can review?
Anyone in the community is free to review the PR once the tests have passed.
cc @zucchini-nlp @vasqu