birder-clip 0.0.2.dev4__tar.gz → 0.0.2.dev6__tar.gz

birder_clip-0.0.2.dev6/PKG-INFO

@@ -0,0 +1,151 @@
+Metadata-Version: 2.4
+Name: birder_clip
+Version: 0.0.2.dev6
+Summary: A Birder extension for CLIP-style image-text modeling and multimodal computer vision workflows.
+Author: Ofer Hasson
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://gitlab.com/birder/birder-clip
+Project-URL: Issues, https://gitlab.com/birder/birder-clip/-/issues
+Project-URL: Changelog, https://gitlab.com/birder/birder-clip/-/blob/main/CHANGELOG.md
+Keywords: computer-vision,clip,image-text,pytorch,deep-learning,artificial intelligence
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: birder>=0.6.0
+Requires-Dist: ftfy>=6.3.1
+Requires-Dist: regex>=2025.7.29
+Requires-Dist: tqdm>=4.67.0
+Requires-Dist: webdataset>=0.2.111
+Requires-Dist: huggingface_hub
+Requires-Dist: transformers
+Requires-Dist: torch>=2.10.0
+Requires-Dist: torchvision
+Provides-Extra: dev
+Requires-Dist: bandit~=1.9.4; extra == "dev"
+Requires-Dist: black~=26.5.0; extra == "dev"
+Requires-Dist: build~=1.5.0; extra == "dev"
+Requires-Dist: bumpver~=2026.1132; extra == "dev"
+Requires-Dist: coverage~=7.14.1; extra == "dev"
+Requires-Dist: debugpy; extra == "dev"
+Requires-Dist: flake8-pep585~=0.1.7; extra == "dev"
+Requires-Dist: flake8~=7.3.0; extra == "dev"
+Requires-Dist: invoke~=3.0.3; extra == "dev"
+Requires-Dist: ipython; extra == "dev"
+Requires-Dist: isort~=8.0.1; extra == "dev"
+Requires-Dist: mkdocs~=1.6.1; extra == "dev"
+Requires-Dist: mkdocs-exclude~=1.0.2; extra == "dev"
+Requires-Dist: mypy~=2.1.0; extra == "dev"
+Requires-Dist: parameterized~=0.9.0; extra == "dev"
+Requires-Dist: pylint~=4.0.6; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: requests~=2.34.2; extra == "dev"
+Requires-Dist: safetensors~=0.7.0; extra == "dev"
+Requires-Dist: setuptools; extra == "dev"
+Requires-Dist: twine~=6.2.0; extra == "dev"
+Requires-Dist: types-requests~=2.33.0; extra == "dev"
+Requires-Dist: urllib3~=2.7.0; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Dynamic: license-file
+
+# Birder CLIP
+
+Birder CLIP is an early-stage Birder extension for CLIP-style image-text models, focused on practical inference and fine-tuning workflows.
+
+- [Introduction](#introduction)
+- [Setup](#setup)
+- [Getting Started](#getting-started)
+- [Training](#training)
+- [Project Status and Contributions](#project-status-and-contributions)
+- [Licenses](#licenses)
+- [Acknowledgments](#acknowledgments)
+
+## Introduction
+
+Birder CLIP extends [Birder](https://gitlab.com/birder/birder) with image-text models for zero-shot classification, image-text retrieval style workflows, caption generation and related multimodal computer vision tasks.
+
+The project is aimed at image-text modeling rather than general vision-language model (VLM) chat or instruction-following systems.
+It currently includes CLIP-style components, tokenizers, model registry utilities, inference scripts and training code.
+Full training is supported, but for large-scale CLIP pretraining you are probably better served by [OpenCLIP](https://github.com/mlfoundations/open_clip).
+
+## Setup
+
+1. Ensure your environment meets the minimum requirements:
+   - Python 3.11 or newer
+   - PyTorch 2.10 or newer (installed for your hardware/driver stack)
+   - Birder 0.6.0 or newer
+
+1. Install the latest Birder CLIP version:
+
+```sh
+pip install birder-clip
+```
+
+## Getting Started
+
+List available image-text models:
+
+```sh
+python -m birder_clip.tools list-models --image-text
+```
+
+List available pretrained weights:
+
+```sh
+python -m birder_clip.tools list-models --pretrained --verbose
+```
+
+Run zero-shot classification on a directory of images:
+
+```sh
+python -m birder_clip.scripts.zero_shot -n laion_clip_vit_l14 --classes eagle hawk falcon --template-set default --gpu data/images
+```
+
+For detailed options, run:
+
+```sh
+python -m birder_clip.scripts.zero_shot --help
+python -m birder_clip.tools --help
+```
+
+## Training
+
+Birder CLIP includes training support for image-text datasets in CSV and WebDataset formats, including CLIP, CoCa and LiT-style workflows.
+
+## Project Status and Contributions
+
+Birder CLIP is an early alpha project. APIs, model names, checkpoints, training recipes and command-line options may change without notice.
+
+This is currently a personal project in active development. Suggestions, bug reports and feedback are welcome through the project's issue tracker, but the project is not yet stable enough for broad external contributions.
+
+## Licenses
+
+The code in this project is primarily licensed under Apache 2.0. See [LICENSE](LICENSE) for details.
+
+Some model implementations, pretrained weights, tokenizers and converted artifacts may be derived from or depend on projects and datasets with their own licenses and usage restrictions.
+
+**You are responsible for ensuring compliance with all licenses and conditions of any dependent licenses.**
+
+### Disclaimer
+
+If you intend to use Birder CLIP, its pretrained weights, or any associated datasets in a commercial product, we strongly recommend seeking legal advice to ensure compliance with all relevant licenses and terms of use.
+
+It's the user's responsibility to ensure that their use of this project, including any pretrained weights or datasets, complies with all applicable licenses and legal requirements.
+
+## Acknowledgments
+
+Birder CLIP owes much to the work of others in computer vision, image-text representation learning and open-source machine learning.
+
+Special thanks to the [OpenCLIP](https://github.com/mlfoundations/open_clip) project, which serves as the main reference implementation and inspiration for much of the CLIP-style modeling and training work here. The same principle as in Birder applies: this project stands on the shoulders of many open-source projects, papers and datasets. If an attribution is missing, please open an issue.

birder_clip-0.0.2.dev6/README.md

@@ -0,0 +1,89 @@
+# Birder CLIP
+
+Birder CLIP is an early-stage Birder extension for CLIP-style image-text models, focused on practical inference and fine-tuning workflows.
+
+- [Introduction](#introduction)
+- [Setup](#setup)
+- [Getting Started](#getting-started)
+- [Training](#training)
+- [Project Status and Contributions](#project-status-and-contributions)
+- [Licenses](#licenses)
+- [Acknowledgments](#acknowledgments)
+
+## Introduction
+
+Birder CLIP extends [Birder](https://gitlab.com/birder/birder) with image-text models for zero-shot classification, image-text retrieval style workflows, caption generation and related multimodal computer vision tasks.
+
+The project is aimed at image-text modeling rather than general vision-language model (VLM) chat or instruction-following systems.
+It currently includes CLIP-style components, tokenizers, model registry utilities, inference scripts and training code.
+Full training is supported, but for large-scale CLIP pretraining you are probably better served by [OpenCLIP](https://github.com/mlfoundations/open_clip).
+
+## Setup
+
+1. Ensure your environment meets the minimum requirements:
+   - Python 3.11 or newer
+   - PyTorch 2.10 or newer (installed for your hardware/driver stack)
+   - Birder 0.6.0 or newer
+
+1. Install the latest Birder CLIP version:
+
+```sh
+pip install birder-clip
+```
+
+## Getting Started
+
+List available image-text models:
+
+```sh
+python -m birder_clip.tools list-models --image-text
+```
+
+List available pretrained weights:
+
+```sh
+python -m birder_clip.tools list-models --pretrained --verbose
+```
+
+Run zero-shot classification on a directory of images:
+
+```sh
+python -m birder_clip.scripts.zero_shot -n laion_clip_vit_l14 --classes eagle hawk falcon --template-set default --gpu data/images
+```
+
+For detailed options, run:
+
+```sh
+python -m birder_clip.scripts.zero_shot --help
+python -m birder_clip.tools --help
+```
+
+## Training
+
+Birder CLIP includes training support for image-text datasets in CSV and WebDataset formats, including CLIP, CoCa and LiT-style workflows.
+
+## Project Status and Contributions
+
+Birder CLIP is an early alpha project. APIs, model names, checkpoints, training recipes and command-line options may change without notice.
+
+This is currently a personal project in active development. Suggestions, bug reports and feedback are welcome through the project's issue tracker, but the project is not yet stable enough for broad external contributions.
+
+## Licenses
+
+The code in this project is primarily licensed under Apache 2.0. See [LICENSE](LICENSE) for details.
+
+Some model implementations, pretrained weights, tokenizers and converted artifacts may be derived from or depend on projects and datasets with their own licenses and usage restrictions.
+
+**You are responsible for ensuring compliance with all licenses and conditions of any dependent licenses.**
+
+### Disclaimer
+
+If you intend to use Birder CLIP, its pretrained weights, or any associated datasets in a commercial product, we strongly recommend seeking legal advice to ensure compliance with all relevant licenses and terms of use.
+
+It's the user's responsibility to ensure that their use of this project, including any pretrained weights or datasets, complies with all applicable licenses and legal requirements.
+
+## Acknowledgments
+
+Birder CLIP owes much to the work of others in computer vision, image-text representation learning and open-source machine learning.
+
+Special thanks to the [OpenCLIP](https://github.com/mlfoundations/open_clip) project, which serves as the main reference implementation and inspiration for much of the CLIP-style modeling and training work here. The same principle as in Birder applies: this project stands on the shoulders of many open-source projects, papers and datasets. If an attribution is missing, please open an issue.

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev6}/birder_clip/common/fs_ops.py

@@ -67,7 +67,7 @@ def model_path(
     network_name: str,
     *,
     epoch: Optional[int | str] = None,
-    file_format: FileFormatType = "pt",
+    st: bool = False,
     states: bool = False,
 ) -> Path:
     if epoch is not None:
@@ -77,8 +77,10 @@ def model_path(
 
     if states is True:
         file_name = f"{file_name}_states.pt"
+    elif st is True:
+        file_name = f"{file_name}.safetensors"
     else:
-        file_name = f"{file_name}.{file_format}"
+        file_name = f"{file_name}.pt"
 
     return settings.MODELS_DIR.joinpath(file_name)
 
@@ -109,6 +111,30 @@ def _checkpoint_states(
     torch.save(kwargs, states_path)
 
 
+def _checkpoint_states_from_state_dicts(
+    states_path: Path,
+    optimizer_state: Optional[dict[str, Any]],
+    scheduler_state: Optional[dict[str, Any]],
+    scaler_state: Optional[dict[str, Any]],
+    model_base_state: Optional[dict[str, Any]],
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    if optimizer_state is None or scheduler_state is None:
+        return
+
+    logger.info(f"Saving checkpoint states {states_path}...")
+    torch.save(
+        {
+            "optimizer_state": optimizer_state,
+            "scheduler_state": scheduler_state,
+            "scaler_state": scaler_state,
+            "model_base_state": model_base_state,
+            **extra_states,
+        },
+        states_path,
+    )
+
+
 class TrainingStates(NamedTuple):
     optimizer_state: Optional[dict[str, Any]]
     scheduler_state: Optional[dict[str, Any]]
@@ -182,6 +208,50 @@ def checkpoint_model(
     _checkpoint_states(states_path, optimizer, scheduler, scaler, model_base, **extra_states)
 
 
+def checkpoint_model_from_state_dicts(
+    network_name: str,
+    epoch: int,
+    model_state: dict[str, Any],
+    task: Any,
+    signature: SignatureType,
+    rgb_stats: RGBType,
+    optimizer_state: Optional[dict[str, Any]],
+    scheduler_state: Optional[dict[str, Any]],
+    scaler_state: Optional[dict[str, Any]],
+    model_base_state: Optional[dict[str, Any]],
+    *,
+    external_config: Optional[dict[str, Any]] = None,
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    kwargs = {}
+    if external_config is not None:
+        kwargs["config"] = external_config
+
+    path = model_path(network_name, epoch=epoch)
+    states_path = model_path(network_name, epoch=epoch, states=True)
+    logger.info(f"Saving model checkpoint {path}...")
+    torch.save(
+        {
+            "state": model_state,
+            "birder_clip_version": __version__,
+            "task": task,
+            "signature": signature,
+            "rgb_stats": rgb_stats,
+            **kwargs,
+        },
+        path,
+    )
+
+    _checkpoint_states_from_state_dicts(
+        states_path,
+        optimizer_state,
+        scheduler_state,
+        scaler_state,
+        model_base_state,
+        **extra_states,
+    )
+
+
 def clean_checkpoints(network_name: str, keep_last: int) -> None:
     epoch = "*[0-9]"
     models_glob = str(model_path(network_name, epoch=epoch))
@@ -314,7 +384,7 @@ def load_model(
             embed_dim=embed_dim,
             tokenizer=tokenizer,
         )
-        path = model_path(_network_name, epoch=epoch, file_format="safetensors" if st is True else "pt")
+        path = model_path(_network_name, epoch=epoch, st=st)
 
     logger.info(f"Loading model from {path} on device {device}...")
 
@@ -589,6 +659,33 @@ def load_pretrained_tokenizer(weights: str, *, download: bool = True, **kwargs:
     return get_tokenizer(tokenizer_name, **tokenizer_kwargs)
 
 
+def save_st(
+    net: torch.nn.Module,
+    dst: str,
+    task: str,
+    signature: SignatureType,
+    rgb_stats: RGBType,
+    *,
+    external_config: Optional[dict[str, Any]] = None,
+) -> None:
+    assert _HAS_SAFETENSORS, "'pip install safetensors' to use .safetensors"
+    kwargs = {}
+    if external_config is not None:
+        kwargs["config"] = json.dumps(external_config)
+
+    safetensors.torch.save_model(
+        net,
+        str(dst),
+        {
+            "birder_clip_version": __version__,
+            "task": task,
+            "signature": json.dumps(signature),
+            "rgb_stats": json.dumps(rgb_stats),
+            **kwargs,
+        },
+    )
+
+
 def download_model_by_weights(
     weights: str, *, dst: Optional[str | Path] = None, file_format: FileFormatType = "pt", progress_bar: bool = True
 ) -> None:

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev6}/birder_clip/common/lib.py

@@ -43,7 +43,7 @@ def get_image_text_network_name(
     parts = [network]
     if image_encoder is not None:
         parts.append(image_encoder)
-    if text_encoder is not None and text_encoder != "text_transformer":
+    if text_encoder is not None and text_encoder != "transformer_encoder":
         parts.append(text_encoder)
 
     if registry.exists(network) is True:
@@ -86,7 +86,7 @@ def get_image_text_model_config(
 
     if config is not None:
         for key, value in config.items():
-            if key in {"image", "text"} and isinstance(value, dict) is True:
+            if key in {"image", "text"} and isinstance(value, dict):
                 model_config[key] = {**model_config.get(key, {}), **value}
             else:
                 model_config[key] = value

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev6}/birder_clip/common/training_cli.py

@@ -21,6 +21,11 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("-n", "--network", type=str, help="the image-text network to train")
     parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument("--image-encoder", type=str, help="the image encoder to use")
+    parser.add_argument(
+        "--image-encoder-pretrained",
+        type=str,
+        help="pretrained Birder image model weights path to load into the image encoder",
+    )
     parser.add_argument("--text-encoder", type=str, help="the text encoder to use")
     parser.add_argument("--embed-dim", type=int, metavar="N", help="shared image-text embedding dimension")
     parser.add_argument("--tokenizer", type=str, help="the tokenizer to use")
@@ -43,13 +48,31 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
 
 def add_loss_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Loss parameters")
-    group.add_argument("--loss", type=str, choices=["clip"], default="clip", help="loss function to use")
+    group.add_argument("--loss", type=str, choices=["clip", "coca"], default="clip", help="loss function to use")
+    group.add_argument(
+        "--coca-caption-loss-weight", type=float, default=1.0, help="weight assigned to CoCa caption loss"
+    )
+    group.add_argument(
+        "--coca-contrastive-loss-weight", type=float, default=1.0, help="weight assigned to CoCa contrastive loss"
+    )
+
+
+def add_freeze_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Freeze parameters")
+    group.add_argument(
+        "--freeze-image-encoder",
+        default=False,
+        action="store_true",
+        help="freeze image encoder body, leaving the projection head trainable",
+    )
 
 
 def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: int = 32) -> None:
     group = parser.add_argument_group("Optimization parameters")
     group.add_argument("--batch-size", type=int, default=default_batch_size, metavar="N", help="the batch size")
-    group.add_argument("--opt", type=str, choices=list(get_args(OptimizerType)), default="sgd", help="optimizer to use")
+    group.add_argument(
+        "--opt", type=str, choices=list(get_args(OptimizerType)), default="adamw", help="optimizer to use"
+    )
     group.add_argument("--opt-fused", default=False, action="store_true", help="use fused optimizer implementation")
     group.add_argument("--momentum", type=float, default=0.9, metavar="M", help="optimizer momentum")
     group.add_argument("--nesterov", default=False, action="store_true", help="use nesterov momentum")
@@ -64,7 +87,12 @@ def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: i
         metavar="N",
         help="number of iterations to accumulate gradients per optimizer step",
     )
-    # NOTE: Add flag for negative sample caching in grad accum mode
+    group.add_argument(
+        "--grad-accum-cache-negatives",
+        default=False,
+        action="store_true",
+        help="cache features so CLIP loss uses all accumulated microbatches as negatives",
+    )
 
 
 def add_lr_wd_args(parser: argparse.ArgumentParser) -> None:
@@ -248,6 +276,8 @@ def add_data_aug_args(
     )
     group.add_argument("--ra-magnitude", type=int, default=9, help="magnitude for all the RandAugment transformations")
     group.add_argument("--augmix-severity", type=int, default=3, help="severity of AugMix policy")
+    group.add_argument("--clip-color-jitter-prob", type=float, default=0.0, help="CLIP color jitter probability")
+    group.add_argument("--clip-gray-prob", type=float, default=0.0, help="CLIP grayscale probability")
     group.add_argument("--resize-min-scale", type=float, default=default_min_scale, help="random resize min scale")
     group.add_argument(
         "--re-prob",
@@ -318,7 +348,13 @@ def add_dataloader_args(parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="keep dataloader worker processes alive between epochs",
     )
-    group.add_argument("--drop-last", default=False, action="store_true", help="drop the last incomplete batch")
+    group.add_argument(
+        "--no-drop-last",
+        dest="drop_last",
+        default=True,
+        action="store_false",
+        help="do not drop the last incomplete batch",
+    )
 
 
 def add_precision_args(parser: argparse.ArgumentParser) -> None:
@@ -348,6 +384,29 @@ def add_precision_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
+def add_grad_checkpointing_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Gradient checkpointing parameters")
+    group.add_argument(
+        "--grad-checkpointing",
+        default=False,
+        action="store_true",
+        help="enable gradient checkpointing for supported models",
+    )
+    group.add_argument(
+        "--grad-checkpointing-segments",
+        type=int,
+        metavar="N",
+        help="number of checkpoint segments to request from supported models",
+    )
+    group.add_argument(
+        "--no-grad-checkpointing-preserve-rng-state",
+        dest="grad_checkpointing_preserve_rng_state",
+        default=True,
+        action="store_false",
+        help="disable RNG state preservation during gradient checkpointing recomputation",
+    )
+
+
 def add_compile_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Compilation parameters")
     group.add_argument("--compile", default=False, action="store_true", help="enable compilation")
@@ -410,6 +469,44 @@ def add_distributed_args(parser: argparse.ArgumentParser) -> None:
     group.add_argument("--local-rank", type=int, metavar="N", help="local rank")
     group.add_argument("--dist-url", type=str, default="env://", help="URL used to initialize distributed training")
     group.add_argument("--dist-backend", type=str, default="nccl", help="distributed backend")
+    group.add_argument(
+        "--distributed-mode", type=str, choices=["ddp", "fsdp"], default="ddp", help="distributed training mode"
+    )
+    group.add_argument(
+        "--fsdp-sharding-strategy",
+        type=str,
+        choices=["shard-grad-op", "full-shard"],
+        default="shard-grad-op",
+        help="FSDP sharding strategy",
+    )
+    group.add_argument(
+        "--fsdp-param-dtype",
+        type=str,
+        choices=["float32", "float16", "bfloat16"],
+        help="FSDP mixed precision parameter dtype",
+    )
+    group.add_argument(
+        "--fsdp-reduce-dtype",
+        type=str,
+        choices=["float32", "float16", "bfloat16"],
+        help="FSDP mixed precision gradient reduction dtype",
+    )
+    group.add_argument(
+        "--fsdp-wrap-policy",
+        type=str,
+        choices=["block-group-regex", "min-num-params"],
+        default="block-group-regex",
+        help="FSDP module wrapping policy",
+    )
+    group.add_argument(
+        "--fsdp-wrap-min-num-params",
+        type=float,
+        metavar="M",
+        help="minimum module parameter count in millions for wrapping when using --fsdp-wrap-policy min-num-params",
+    )
+    group.add_argument(
+        "--fsdp-offload-policy", type=str, choices=["none", "cpu"], default="none", help="FSDP parameter offload policy"
+    )
     group.add_argument(
         "--find-unused-parameters",
         default=False,
@@ -532,8 +629,12 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--load-states requires --resume-epoch to be set")
     if args.load_scheduler is True and args.resume_epoch is None:
         raise cli.ValidationError("--load-scheduler requires --resume-epoch to be set")
-    if hasattr(args, "pretrained") is True and args.pretrained is True and args.resume_epoch is not None:
+    if args.pretrained is True and args.resume_epoch is not None:
         raise cli.ValidationError("--pretrained cannot be used with --resume-epoch")
+    if args.image_encoder_pretrained is not None and args.resume_epoch is not None:
+        raise cli.ValidationError("--image-encoder-pretrained cannot be used with --resume-epoch")
+    if args.pretrained is True and args.image_encoder_pretrained is not None:
+        raise cli.ValidationError("--image-encoder-pretrained cannot be used with --pretrained")
 
     if args.freeze_bn is True and args.sync_bn is True:
         raise cli.ValidationError("--freeze-bn cannot be used with --sync-bn")
@@ -559,5 +660,44 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--context-length must be positive")
     if args.grad_accum_steps < 1:
         raise cli.ValidationError("--grad-accum-steps must be >= 1")
+    if args.grad_accum_cache_negatives is True and args.grad_accum_steps == 1:
+        raise cli.ValidationError("--grad-accum-cache-negatives requires --grad-accum-steps greater than 1")
+    if args.grad_accum_cache_negatives is True and args.loss == "coca":
+        raise cli.ValidationError("--grad-accum-cache-negatives is only supported with --loss clip")
+
+    if args.coca_caption_loss_weight < 0.0:
+        raise cli.ValidationError("--coca-caption-loss-weight must be non-negative")
+    if args.coca_contrastive_loss_weight < 0.0:
+        raise cli.ValidationError("--coca-contrastive-loss-weight must be non-negative")
+
+    # EMA
     if args.model_ema_steps < 1:
         raise cli.ValidationError("--model-ema-steps must be >= 1")
+
+    # Gradient checkpointing args
+    if args.grad_checkpointing_segments is not None and args.grad_checkpointing_segments < 1:
+        raise cli.ValidationError("--grad-checkpointing-segments must be >= 1")
+    if args.grad_checkpointing_segments is not None and args.grad_checkpointing is False:
+        raise cli.ValidationError("--grad-checkpointing-segments requires --grad-checkpointing")
+
+    if args.distributed_mode == "fsdp":
+        if args.grad_checkpointing is True:
+            raise cli.ValidationError("--grad-checkpointing cannot be used with --distributed-mode fsdp")
+        if args.sync_bn is True:
+            raise cli.ValidationError("--sync-bn cannot be used with --distributed-mode fsdp")
+        if args.find_unused_parameters is True:
+            raise cli.ValidationError("--find-unused-parameters cannot be used with --distributed-mode fsdp")
+        if args.compile_opt is True:
+            raise cli.ValidationError("--compile-opt cannot be used with --distributed-mode fsdp")
+        if args.compile_fullgraph is True:
+            raise cli.ValidationError("--compile-fullgraph cannot be used with --distributed-mode fsdp")
+        if args.cpu is True:
+            raise cli.ValidationError("--cpu cannot be used with --distributed-mode fsdp")
+        if args.model_ema is True:
+            raise cli.ValidationError("--model-ema cannot be used with --distributed-mode fsdp")
+        if args.fsdp_wrap_policy == "min-num-params" and args.fsdp_wrap_min_num_params is None:
+            raise cli.ValidationError(
+                "--fsdp-wrap-min-num-params is required when --fsdp-wrap-policy is min-num-params"
+            )
+        if args.fsdp_wrap_min_num_params is not None and args.fsdp_wrap_min_num_params <= 0:
+            raise cli.ValidationError("--fsdp-wrap-min-num-params must be > 0")