open_clip

OpenCLIP

[Paper] [Citations] [Clip Colab] [Coca Colab]

⚠️ Main branch training stack notice

main now uses the post-refactor training stack by default. Training is organized around TrainingTask wrappers, dict-based batches, FSDP2 support, NaFlex image/audio pipelines, and multiple torch.compile strategies. The scope has grown well beyond the original refactor and now includes several new model families. For the older release-stable training API, pin to the v3 branch or the latest 3.x release on PyPI. Inference usage of pretrained image/text models is still intended to be compatible, but training scripts and downstream integrations should review the changes below before upgrading.

New / experimental model families on main:

• NaFlex CLIP — variable-resolution/aspect image towers (timm naflexvit) with token-budget batching (--use-naflex, naflex_* configs)
• NaFlex CLAP — audio-text contrastive training with variable-duration audio (naflexclap_* configs)
• NaFlex GenLIP / GenLAP — generative image/audio captioning with prefix-LM attention and packed [media ; text] rows (naflexgenlip_*, naflexgenlap_* configs, tiktoken text)
• Modern text tower — text_cfg.text_arch="modern": RoPE, SwiGLU/ReLU², RMSNorm, masked pooling (eos/mean/map), with optional qk-norm, gated attention, register tokens, sandwich norm, value residuals, and zero-init residual (moderntext-* configs)
• Variable-length text — text_cfg.variable_text=true pads captions to the per-batch max instead of a fixed context length (works with the modern tower and HF towers)
• Hugging Face ModernBERT text towers — e.g. gte-modernbert-base-ViT-B-32-256
• MaMMUT — single text decoder used in two passes per the paper: bi-directional without cross-attention for contrastive learning, causal with cross-attention for captioning (trains via the CoCa task/loss). Two config families, mirroring coca/coca2: mammut_* reproduces the original LAION-fork numerics exactly via legacy flags (pool_type="avg_all", use_pad_mask=false) and loads the released LAION openMaMMUT-ViT-L-14 weights (pretrained tag datacomp1b_s12_8b_b180k, or directly via hf-hub: — fork-format configs and state dicts are translated on load); mammut2_* uses the corrected defaults (masked-mean text pooling, pad masking). A modern-arch decoder variant is available via multimodal_cfg.text_arch="modern" (mammut2-moderntext_*)
• CoCa v2 configs — coca2_*: paper-faithful attentional pooling (vision_cfg.attentional_pool="cascade", the paper's default — a separate single-query contrastive pooler, so the caption decoder now cross-attends over all 256 generative pooler tokens instead of 255, fixing #458) and the corrected CLS/pad attention mask (text_cfg.correct_cls_mask=true), per #554. A modern-text variant pairs the modern tower with a modern multimodal decoder (coca2-moderntext_*). Existing coca_* configs and released weights are unchanged
• Text validity masks for generative models — CoCa/MaMMUT forward()/encode_text() accept text_valid ([B, L], True = real token); tokenizers can emit exact masks (tokenizer(texts, output_mask=True)), and caption labels are masked to -100 from them. Fixes the SimpleTokenizer pad-collision class (id 0 is a real token, '!' merges like x!=y emit it mid-caption); absent a mask, behavior falls back to the historical text != pad_id derivation. Text towers keep the HF-style attention_mask kwarg at the tower boundary
• Variable text with gradient accumulation — CoCa/MaMMUT support different caption lengths across microbatches in both training entrypoints. Model forwards retain each microbatch's text length; caption logits and masked labels are padded when combining loss inputs, preserving the mean over all valid target tokens. This applies to the logits-based caption loss; --fused-caption-loss still requires --accum-freq 1.

Breaking changes — training CLI:

• --horovod removed (Horovod support deleted; DDP/FSDP2 only)
• --torchscript and --trace removed (torch.jit is being deprecated upstream)
• Default --precision changed from amp → amp_bf16 (silent behavior change; pass --precision amp explicitly to keep fp16 AMP)
• SigLIP's --loss-dist-impl now defaults to gather, as does standalone SigLipLoss. Pass --loss-dist-impl bidir to keep bidirectional ring exchange; reduce and shift remain available. Gather stores all ranks' text features on each rank.
• --naflex-max-tokens-per-batch now defaults to unset. The local token budget is inferred as --batch-size * max(--naflex-seq-lens); GenLIP/GenLAP also include their caption-token cap in the per-row cost. Pass an explicit token budget to preserve older runs that relied on the previous 16384 default.

New training CLI flags (opt-in):

• --siglip-chunk-size — image rows per SigLIP logits chunk (0 disables). For example, --siglip --siglip-chunk-size 1024 enables chunking in both current and legacy training.
• --fsdp — use FSDP2 (fully_shard) instead of DDP
• --fsdp-no-reshard-after-forward, --fsdp-offload-cpu
• --fsdp-checkpoint {full,sharded} — full gathers to rank-0 as a single .pt; sharded uses DCP per-rank shards (faster, lower memory)
• --torchcompile-strategy {task,model,step} — choose whether torch.compile captures task forward/loss, the underlying model, or the full single-batch train step
• --use-naflex and --naflex-* flags — enable NaFlex variable-aspect image pipelines for compatible timm/OpenCLIP ViT-family models (token-budget batching via --naflex-seq-lens / --naflex-max-tokens-per-batch; the same machinery drives the NaFlex audio and generative models)
• --audio-* and --audio-zeroshot-* flags — enable CLAP audio preprocessing/training and Hugging Face audio zero-shot evaluation
• --length-bucketing, --bucket-pool, --bucket-chunk — reorder the train stream by sample length (caption and/or audio tokens) to tighten per-batch padding; the bucket pool holds raw, undecoded samples
• --text-pad-multiple — round per-batch variable-text length up to a multiple, bounding the distinct sequence lengths torch.compile sees (text-axis analogue of --naflex-pad-multiple)
• --text-attention-mask — emit a per-sample text validity mask (batch key text_valid) from the tokenizer, consumed by CoCa/MaMMUT for attention/pooling and -100 caption-label masking. Default auto-enables for CoCa/MaMMUT (except under --distill) and is rejected for tasks that don't consume it
• --caption-z-loss-weight, --caption-loss-compute-dtype {float32,model}, and --caption-loss-chunk-size — configure the next-token objective shared by CoCa/MaMMUT and GenLIP/GenLAP. Defaults preserve the existing fp32 CE with no z-loss; model preserves the loss-logit dtype and ambient autocast policy while returned loss/component scalars remain fp32

Breaking changes — Python API:

• trace_model removed from the top-level open_clip namespace
• load_openai_model, list_openai_models, and build_model_from_openai_state_dict removed. Original-OpenAI weights are still loadable through the standard create_model_from_pretrained(..., pretrained='openai') path, which now routes through HuggingFace Hub (timm/*_clip.openai) instead of torch.jit.load on openaipublic.azureedge.net archives. Removing the JIT path closes an arbitrary-code-execution surface (JIT archives can ship pickled code).
• Training pipeline wraps model + loss in a TrainingTask subclass (CLIPTask, SigLIPTask, CoCaTask, DistillCLIPTask, CLAPTask). Code that previously called train_one_epoch(model, loss, ...) or evaluate(model, ...) directly should switch to passing a task. Tasks construct their own losses; create_loss("clip", ...) is available for standalone training loops (see below).
• Data pipelines emit dict batches instead of tuples. Image/text loaders use {"image": ..., "text": ...}; CLAP audio loaders use {"audio": ..., "text": ...}. Tuple-style image/text calls through task(images, texts) still work via a backward-compat path, but downstream code that iterates a dataloader directly should read the named keys.
• CoCa's autoregressive label shift moved out of coca_model.py and into CoCaTask. coca_model.forward() no longer performs the [:, :-1] / [:, 1:] slicing — callers that relied on that behavior outside training should handle labels themselves.
• CoCa's exact-mask API adds text_valid after text in encode_text, forward, and forward_intermediates. This shifts the older trailing positional arguments (normalize, image_latent, image_indices, and so on); pass those arguments by keyword. For example, replace model.encode_text(text, False) with model.encode_text(text, normalize=False).
• CLIPTextCfg.eos_id no longer defaults to 2 (that value is only correct for XLM-style vocabs). Configs using pool_type="eos" must set eos_id explicitly, and get_tokenizer now validates eos_id/pad_id against the resolved tokenizer, raising on mismatch instead of pooling/masking silently wrong positions.
• HFTokenizer no longer fabricates pad_token_id=0 when the underlying tokenizer has no pad token (id 0 is a real token in most BPE vocabs); variable-text setups fail fast instead. It also forces padding_side='right', which all OpenCLIP pooling/masking assumes.
• Tokenizer wrappers now share special-token controls: encode(..., add_special_tokens=False) remains body-only by default, while model-facing tokenizer(...) defaults to add_special_tokens=True. decode() / batch_decode() default to skip_special_tokens=False, stop_at_eos=True; pass stop_at_eos=False to inspect tokens after the first EOS. This intentionally changes legacy SimpleTokenizer decode output by hiding post-EOS id-0 fill (!) and makes TikToken decode render its reserved control tokens unless skip_special_tokens=True.
• Model traits replace model-name sniffing. Every model carries model.traits (open_clip.get_model_traits(model)): family, objectives, NaFlex/variable-text contracts. The training entrypoints derive NaFlex data, variable text, the --text-attention-mask default and the grad-accum / distill guards from the built model via apply_model_traits, so hf-hub: and renamed configs no longer need a magic substring; args.genlip / args.genlap / args.naflexclap are gone and the data loaders (get_data, get_wds_dataset, ...) take model_traits. --use-naflex now always sets force_naflex_vision, which the factory treats as a no-op for NaFlex-native and audio models. GenLIP configs get NaFlex transforms from a plain create_model_and_transforms call. create_task() likewise selects the task from the built model. MaMMUT decoders now honor multimodal_cfg.variable_text (previously dropped), so mammut2-moderntext_* configs train with per-batch padded text.
• MaxPooler (hf_pooler_type="max_pooler") mask polarity fixed — it previously max-pooled over the padding positions instead of the valid ones.
• CoCa.__init__ no longer takes a pad_id argument — model.pad_id is derived from the text tower (the id it actually masks with: text_cfg.pad_id for native towers, the transformers config pad for HF towers). MaMMUT follows the sam