Granite for Docling - #49012
Granite for Docling#49012nassarofficial wants to merge 3 commits into
Conversation
Add the GraniteForDocling vision-language model for document conversion. Pages are tiled at 512x512 with a pixel-shuffle connector, DeepStack injection of intermediate vision features, and a dense Granite-style text decoder that generates DocLang. The same modeling code loads different GraniteForDocling sizes and configurations via text_config and vision_config, and different projectors (high-resolution or density router). The fine path can be selected per request with fine_route=True; checkpoints trained with the coarse path only set use_fine_route=False and do not build it. Optional multi-token prediction heads (num_mtp_layers) add an auxiliary training loss. Serving engines such as vLLM can use the same heads as draft heads for speculative decoding; generate does not. The model is written as a modular file; the configuration, image processors (torchvision and PIL backends), processor and modeling files are generated from it. The image processors tile pages on the grid closest to their aspect ratio, capped at 16 tiles per side to match the tokenizer's tile position markers, and pad batches on the tile dimension. Co-authored-by: Matteo Omenetti <omenetti.matteo@gmail.com>
Keep multi-token prediction spans based on the original sequence length so later heads retain valid positions. Preserve projected DeepStack features when reusing image embeddings, including cached and beam generation. Resolve the fake image token ID after special-token registration and preserve DocLang markup in the pipeline example. Regenerate the standalone model and processor from the modular source and add regression coverage. The per-call image_seq_len override remains a documented known issue.
zucchini-nlp
left a comment
zucchini-nlp
left a comment
There was a problem hiding this comment.
i dont see a reason to ship a new model while Docling is already supported without any custom code, it maps to existing core classes
Thanks for replying, this is for a new model, the new checkpoints are a different architecture, its not idefics3. |
|
Is that a not-yet-released checkpoint, since the link from PR description gives me 404...? (if yes, are you on slack? Just so we can sort out planned release dates and review timely) |
| - Set `padding_side="left"` before batched generation, otherwise the padded prompts generate from the wrong position. | ||
|
|
||
| ```py | ||
| processor.tokenizer.padding_side = "left" | ||
| ``` | ||
|
|
There was a problem hiding this comment.
can be saved in config as default imo, but overall the info is not new and needn't be in docs
| dtype=torch.bfloat16, | ||
| device_map="auto", | ||
| attn_implementation="sdpa", # or "flash_attention_2" |
There was a problem hiding this comment.
dtype can be omitted (default dtype is used from config) and same for attn_implementation (sdpa is the default already)
Review feedback: the config already sets the dtype and sdpa is the default; padding_side=left ships in the checkpoint's tokenizer config instead.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, granite_for_docling |
CI recapDashboard: View test results in Grafana |
There was a problem hiding this comment.
Did a very quick and early review, very good usage of modular!
The main points are:
- MTP which should be re-routed via an existing core API (ping Cyril for help in slack)
- A bit of re-ordering needed in vision model and image processors, mostly just moving code around and i think we could copy from qwen3-vl when it comes to deepstack features. Maybe even inherit/reuse from
Granite4Visionsince you have non-contiguous deepstack layer ids i guess
| "Kwargs passed to `processor.__call__` have to be in `processor_kwargs` dict, not in `**kwargs`" | ||
| ) | ||
| processor_kwargs = processor_kwargs_from_kwargs | ||
| processor_kwargs = {**processor_kwargs, **processor_kwargs_from_kwargs} |
There was a problem hiding this comment.
hmmm, i would really not want users pass processor_kwargs as dict while still passing some kwargs directly
| @lru_cache(maxsize=10) | ||
| def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> list[tuple[int, int]]: | ||
| """ | ||
| Computes all `(num_columns, num_rows)` tile grids holding between `min_image_tiles` and `max_image_tiles` tiles, | ||
| with at most `MAX_TILES_PER_SIDE` tiles per side. | ||
| """ | ||
| max_tiles_per_side = min(max_image_tiles, MAX_TILES_PER_SIDE) | ||
| aspect_ratios = [ | ||
| (width, height) | ||
| for width in range(1, max_tiles_per_side + 1) | ||
| for height in range(1, max_tiles_per_side + 1) | ||
| if min_image_tiles <= width * height <= max_image_tiles | ||
| ] | ||
| return sorted(aspect_ratios, key=lambda x: x[0] * x[1]) |
There was a problem hiding this comment.
could be imported/copied from GotOCR's and we move max_tiles_per_side = min(max_image_tiles, MAX_TILES_PER_SIDE) outside the fn
Also, not clear why we need a max bound for max_tiles, i.e. in which situations the fn gets an exceeding max tiles and needs to be capped?
| @lru_cache(maxsize=100) | ||
| def get_optimal_tiled_canvas( | ||
| original_image_size: tuple[int, int], |
There was a problem hiding this comment.
identical to GotOCR, should be just imported (copied with modular)
| for sample in images: | ||
| sample_tiles, sample_rows, sample_cols = [], [], [] | ||
| for image in sample: | ||
| if crop_to_patches: | ||
| num_cols, num_rows = get_optimal_tiled_canvas( | ||
| tuple(image.shape[-2:]), (size.height, size.width), min_patches, max_patches | ||
| ) | ||
| tiles = self.crop_image_to_patches( | ||
| image[None], min_patches, max_patches, patch_size=size, resample=resample | ||
| )[0] | ||
| else: | ||
| num_cols = num_rows = 1 | ||
| tiles = self.resize(image, size, resample=resample)[None] | ||
| sample_tiles.append(tiles) | ||
| sample_rows.append(num_rows) | ||
| sample_cols.append(num_cols) | ||
| # A text-only sample of the batch has no tiles |
There was a problem hiding this comment.
we should really try to use group_images_by_shape so we don;t loop over each image separately
| # Pad the samples to the same number of tiles with all-zero tiles, which the model discards. | ||
| max_num_tiles = max(len(tiles) for tiles in pixel_values) | ||
| first_tiles = next(tiles for tiles in pixel_values if len(tiles) > 0) | ||
| padded_pixel_values = torch.zeros( | ||
| len(pixel_values), | ||
| max_num_tiles, | ||
| *first_tiles.shape[1:], | ||
| dtype=first_tiles.dtype, | ||
| device=first_tiles.device, | ||
| ) | ||
| tile_fine_mask = torch.zeros(len(pixel_values), max_num_tiles, dtype=torch.bool) | ||
| for i, tiles in enumerate(pixel_values): | ||
| if len(tiles) > 0: | ||
| padded_pixel_values[i, : tiles.shape[0]] = tiles | ||
| tile_fine_mask[i, : tiles.shape[0]] = fine_route |
There was a problem hiding this comment.
smth more like this would do better, maybe we can even use the base class' padding without overriding:
if do_pad:
images = self.pad(images)
| image_outputs = self.vision_model(pixel_values=pixel_values, return_dict=True, **kwargs) | ||
| image_hidden_states = image_outputs.last_hidden_state | ||
| image_features = self.connector(image_hidden_states, tile_fine_mask) | ||
| # `hidden_states[0]` is the patch embedding output, so the output of vision layer `i` is `hidden_states[i + 1]` | ||
| deepstack_features = [ | ||
| self.connector.deepstack(slot, image_outputs.hidden_states[vision_layer_idx + 1], tile_fine_mask) | ||
| for slot, vision_layer_idx in enumerate(self.config.deepstack_visual_indexes) | ||
| ] |
There was a problem hiding this comment.
yeah, we defi can put the Connector and DeepstackModule inside a VisionModel, so that the final returned output has pooler_outputs and depstack_embeds
| if self.density_router is not None: | ||
| tile_sample_index = torch.arange(batch_size, device=pixel_values.device).repeat_interleave(num_tiles) | ||
| router_logits = self.density_router(image_hidden_states, tile_sample_index[real_images_inds], batch_size) | ||
| return GraniteForDoclingImageFeaturesOutput( |
There was a problem hiding this comment.
just to make sure, we get different ckpt with and without routers?
| inputs_embeds: torch.FloatTensor | None = None, | ||
| pixel_values: torch.FloatTensor | None = None, | ||
| tile_fine_mask: torch.BoolTensor | None = None, | ||
| image_hidden_states: torch.FloatTensor | None = None, |
There was a problem hiding this comment.
can we drop image_hidden_states from inputs for now pls, i am doing a major refactor for proper pre-encoded vision states (#45783) and it would be hard to support BC
| if self.training and self.text_model.gradient_checkpointing and use_cache: | ||
| logger.warning_once( | ||
| "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." | ||
| ) | ||
| use_cache = False |
There was a problem hiding this comment.
not needed, it is raised and turned off via decorators
| def prepare_inputs_for_generation(self, *args, is_first_iteration=False, **kwargs): | ||
| model_inputs = super().prepare_inputs_for_generation(*args, is_first_iteration=is_first_iteration, **kwargs) | ||
| if not is_first_iteration and model_inputs.get("past_key_values") is not None: | ||
| model_inputs.pop("image_hidden_states", None) | ||
| model_inputs.pop("deepstack_image_features", None) | ||
| return model_inputs | ||
|
|
There was a problem hiding this comment.
not needed, we can add deepstack_image_features inside a prebuilt set in generation/utils.py instead
This branch has not been deployed
3 participants
What does this PR do?
Adds GraniteForDocling, Docling's IBM Granite vision-language model for document conversion used by
Docling. Given a page image it generates
DocLang: layout elements in reading order with bounding boxes, text, tables,
formulas and code, which
docling-coreturns into Markdown/HTML/JSON.Checkpoints: https://huggingface.co/docling-project/granite-for-docling-500m
().
Architecture (all in
modular_granite_for_docling.py, inheriting from Idefics3 / Granite / GotOcr2):plus a thumbnail when there is more than one tile; a single-tile page gets no thumbnail.
to the image-token positions after the first decoder layers.
tile,
use_fine_route, selected per call withfine_route=True).density_router_hidden_size) that predicts from the encoder featureswhether a page needs the fine path (
model.predict_fine_route).num_mtp_layers) that add an auxiliary training loss;generatedoes not use them (serving engines such as vLLM use them for speculative decoding).Which optional modules a checkpoint has is declared in
config.json, so the same code loads the500M (router + fine path) and larger (MTP) variants.
Tests
tests/models/granite_for_docling/: modeling (incl. common tests, DeepStack tap placement,coarse-only checkpoints, backend equivalence, MTP/router losses, feature reuse), image processing,
processing. Slow integration tests run against the Hub checkpoint on
docling_technical_report_p1.png.Verified against the reference implementation: HF outputs match the training model and the vLLM
implementation token-for-token on a page set covering single-tile, 16x1 strip, coarse and fine routes.
make check-repo/ modular conversion / ruff clean.AI assistance disclosure: parts of the modeling code, tests and docs were drafted with an AI coding
assistant; code was reviewed, edited and tested by the authors, who maintain this model.
Before submitting
This PR fixes a typo or improves the docs
Did you read the contributor guideline and the Pull Request checks?
Was this discussed/approved via a Github issue or the forum? <link if you opened one, else leave unchecked>
Did you make sure to update the documentation with your changes?
Did you write any new necessary tests?
Who can review?
@zucchini-nlp (multimodal / generate) @molbap (vision)