birder-clip 0.0.2.dev3__tar.gz → 0.0.2.dev5__tar.gz

{birder_clip-0.0.2.dev3 → birder_clip-0.0.2.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder_clip
-Version: 0.0.2.dev3
+Version: 0.0.2.dev5
 Summary: A Birder extension for CLIP-style image-text modeling and multimodal computer vision workflows.
 Author: Ofer Hasson
 License-Expression: Apache-2.0
@@ -24,7 +24,7 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: birder>=0.5.2
+Requires-Dist: birder>=0.5.6
 Requires-Dist: ftfy>=6.3.1
 Requires-Dist: regex>=2025.7.29
 Requires-Dist: tqdm>=4.67.0
@@ -37,7 +37,7 @@ 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~=2025.1131; 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"

{birder_clip-0.0.2.dev3 → birder_clip-0.0.2.dev5}/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.dev3 → birder_clip-0.0.2.dev5}/birder_clip/common/lib.py

@@ -43,9 +43,17 @@ def get_image_text_network_name(
     parts = [network]
     if image_encoder is not None:
         parts.append(image_encoder)
-    if text_encoder is not None:
+    if text_encoder is not None and text_encoder != "text_transformer":
         parts.append(text_encoder)
-    if tokenizer is not None:
+
+    if registry.exists(network) is True:
+        default_tokenizer = registry.get_default_tokenizer(network)
+    else:
+        default_tokenizer = "simple_tokenizer"
+    if default_tokenizer is None:
+        default_tokenizer = "simple_tokenizer"
+
+    if tokenizer is not None and tokenizer != default_tokenizer:
         parts.append(tokenizer)
     if embed_dim is not None:
         parts.append(f"d{embed_dim}")
@@ -78,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.dev3 → birder_clip-0.0.2.dev5}/birder_clip/common/training_cli.py

@@ -17,32 +17,6 @@ from birder_clip.model_registry import Task
 from birder_clip.model_registry import registry
 
 
-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")
-    group.add_argument("--compile-fullgraph", default=False, action="store_true", help="compile using fullgraph=True")
-    group.add_argument(
-        "--compile-mode", type=str, choices=list(torch._inductor.list_mode_options().keys()), help="torch.compile mode"
-    )
-    group.add_argument(
-        "--compile-opt", default=False, action="store_true", help="enable compilation for optimizer step"
-    )
-    group.add_argument(
-        "--compile-recompile-limit",
-        type=int,
-        default=torch.compiler.config.recompile_limit,
-        metavar="N",
-        help="maximum recompilations per compiled function before eager fallback",
-    )
-    group.add_argument(
-        "--compile-accumulated-recompile-limit",
-        type=int,
-        default=torch.compiler.config.accumulated_recompile_limit,
-        metavar="N",
-        help="maximum total recompilations across compiled functions",
-    )
-
-
 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")
@@ -92,6 +66,7 @@ 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
 
 
 def add_lr_wd_args(parser: argparse.ArgumentParser) -> None:
@@ -129,14 +104,14 @@ def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
         "--lr-scheduler-update",
         type=str,
         choices=["epoch", "step"],
-        default="step",
+        default="epoch",
         help="when to apply learning rate scheduler update: epoch (once per epoch), step (each optimizer step)",
     )
     group.add_argument(
         "--lr-scheduler",
         type=str,
         choices=list(get_args(SchedulerType)),
-        default="cosine",
+        default="constant",
         help="learning rate scheduler",
     )
     group.add_argument(
@@ -175,15 +150,6 @@ def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
-def add_input_args(parser: argparse.ArgumentParser) -> None:
-    group = parser.add_argument_group("Input parameters")
-    group.add_argument(
-        "--channels", type=int, default=settings.DEFAULT_NUM_CHANNELS, metavar="N", help="no. of image channels"
-    )
-    group.add_argument("--size", type=int, nargs="+", metavar=("H", "W"), help="image size")
-    group.add_argument("--context-length", type=int, metavar="N", help="text context length")
-
-
 def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs: int = 100) -> None:
     group = parser.add_argument_group("Training schedule parameters")
     group.add_argument("--epochs", type=int, default=default_epochs, metavar="N", help="number of training epochs")
@@ -204,6 +170,37 @@ def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs:
     )
 
 
+def add_ema_args(
+    parser: argparse.ArgumentParser, default_ema_steps: int = 1, default_ema_decay: float = 0.9999
+) -> None:
+    group = parser.add_argument_group("Exponential moving average parameters")
+    group.add_argument(
+        "--model-ema",
+        default=False,
+        action="store_true",
+        help="enable tracking exponential moving average of model parameters",
+    )
+    group.add_argument(
+        "--model-ema-steps",
+        type=int,
+        default=default_ema_steps,
+        metavar="N",
+        help="number of optimizer steps between EMA updates",
+    )
+    group.add_argument(
+        "--model-ema-decay",
+        type=float,
+        default=default_ema_decay,
+        help="decay factor for exponential moving average of model parameters",
+    )
+    group.add_argument(
+        "--model-ema-warmup",
+        type=int,
+        metavar="N",
+        help="number of epochs/steps before EMA is applied (defaults to warmup epochs/steps, pass 0 to disable warmup)",
+    )
+
+
 def add_batch_norm_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Batch normalization parameters")
     group.add_argument(
@@ -215,6 +212,15 @@ def add_batch_norm_args(parser: argparse.ArgumentParser) -> None:
     group.add_argument("--sync-bn", default=False, action="store_true", help="use synchronized BatchNorm")
 
 
+def add_input_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Input parameters")
+    group.add_argument(
+        "--channels", type=int, default=settings.DEFAULT_NUM_CHANNELS, metavar="N", help="no. of image channels"
+    )
+    group.add_argument("--size", type=int, nargs="+", metavar=("H", "W"), help="image size")
+    group.add_argument("--context-length", type=int, metavar="N", help="text context length")
+
+
 def add_data_aug_args(
     parser: argparse.ArgumentParser,
     default_level: int = 4,
@@ -260,7 +266,7 @@ def add_data_aug_args(
         "--rgb-mode",
         type=str,
         choices=list(typing.get_args(RGBMode)),
-        default="clip",
+        default="birder",
         help="RGB mean and std to use for normalization",
     )
     group.add_argument(
@@ -279,67 +285,6 @@ def add_data_aug_args(
     )
 
 
-def add_checkpoint_args(parser: argparse.ArgumentParser, default_save_frequency: int = 1) -> None:
-    group = parser.add_argument_group("Checkpoint parameters")
-    group.add_argument(
-        "--save-frequency", type=int, default=default_save_frequency, metavar="N", help="frequency of model saving"
-    )
-    group.add_argument(
-        "--keep-last", type=int, metavar="N", help="number of recent checkpoints to keep (older ones are deleted)"
-    )
-    group.add_argument(
-        "--pretrained",
-        default=False,
-        action="store_true",
-        help="start with pretrained version of specified network (will download if not found locally)",
-    )
-    group.add_argument("--resume-epoch", type=int, metavar="N", help="epoch number to resume training from")
-    group.add_argument(
-        "--non-strict-weights",
-        default=False,
-        action="store_true",
-        help="allow non-strict loading of model weights (missing or unexpected keys in state_dict)",
-    )
-    group.add_argument(
-        "--load-states",
-        default=False,
-        action="store_true",
-        help="load optimizer, scheduler and scaler states when resuming",
-    )
-    group.add_argument("--load-scheduler", default=False, action="store_true", help="load only scheduler when resuming")
-
-
-def add_ema_args(
-    parser: argparse.ArgumentParser, default_ema_steps: int = 1, default_ema_decay: float = 0.9999
-) -> None:
-    group = parser.add_argument_group("Exponential moving average parameters")
-    group.add_argument(
-        "--model-ema",
-        default=False,
-        action="store_true",
-        help="enable tracking exponential moving average of model parameters",
-    )
-    group.add_argument(
-        "--model-ema-steps",
-        type=int,
-        default=default_ema_steps,
-        metavar="N",
-        help="number of optimizer steps between EMA updates",
-    )
-    group.add_argument(
-        "--model-ema-decay",
-        type=float,
-        default=default_ema_decay,
-        help="decay factor for exponential moving average of model parameters",
-    )
-    group.add_argument(
-        "--model-ema-warmup",
-        type=int,
-        metavar="N",
-        help="number of epochs/steps before EMA is applied (defaults to warmup epochs/steps, pass 0 to disable warmup)",
-    )
-
-
 def add_dataloader_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Dataloader parameters")
     group.add_argument(
@@ -375,7 +320,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:
@@ -405,12 +356,106 @@ def add_precision_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
+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")
+    group.add_argument("--compile-fullgraph", default=False, action="store_true", help="compile using fullgraph=True")
+    group.add_argument(
+        "--compile-mode", type=str, choices=list(torch._inductor.list_mode_options().keys()), help="torch.compile mode"
+    )
+    group.add_argument(
+        "--compile-opt", default=False, action="store_true", help="enable compilation for optimizer step"
+    )
+    group.add_argument(
+        "--compile-recompile-limit",
+        type=int,
+        default=torch.compiler.config.recompile_limit,
+        metavar="N",
+        help="maximum recompilations per compiled function before eager fallback",
+    )
+    group.add_argument(
+        "--compile-accumulated-recompile-limit",
+        type=int,
+        default=torch.compiler.config.accumulated_recompile_limit,
+        metavar="N",
+        help="maximum total recompilations across compiled functions",
+    )
+
+
+def add_checkpoint_args(parser: argparse.ArgumentParser, default_save_frequency: int = 1) -> None:
+    group = parser.add_argument_group("Checkpoint parameters")
+    group.add_argument(
+        "--save-frequency", type=int, default=default_save_frequency, metavar="N", help="frequency of model saving"
+    )
+    group.add_argument(
+        "--keep-last", type=int, metavar="N", help="number of recent checkpoints to keep (older ones are deleted)"
+    )
+    group.add_argument(
+        "--pretrained",
+        default=False,
+        action="store_true",
+        help="start with pretrained version of specified network (will download if not found locally)",
+    )
+    group.add_argument("--resume-epoch", type=int, metavar="N", help="epoch number to resume training from")
+    group.add_argument(
+        "--non-strict-weights",
+        default=False,
+        action="store_true",
+        help="allow non-strict loading of model weights (missing or unexpected keys in state_dict)",
+    )
+    group.add_argument(
+        "--load-states",
+        default=False,
+        action="store_true",
+        help="load optimizer, scheduler and scaler states when resuming",
+    )
+    group.add_argument("--load-scheduler", default=False, action="store_true", help="load only scheduler when resuming")
+
+
 def add_distributed_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Distributed training parameters")
     group.add_argument("--world-size", type=int, default=1, metavar="N", help="number of distributed processes")
     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,
@@ -558,5 +603,27 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--embed-dim must be positive")
     if args.context_length is not None and args.context_length <= 0:
         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.model_ema_steps < 1:
         raise cli.ValidationError("--model-ema-steps must be >= 1")
+
+    if args.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")

birder_clip-0.0.2.dev5/birder_clip/common/training_utils.py

@@ -0,0 +1,99 @@
+import argparse
+import logging
+from pathlib import Path
+from typing import Any
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from birder.common import fsdp_utils
+from birder.common import training_utils as birder_training_utils
+
+from birder_clip.common import fs_ops
+from birder_clip.conf import settings
+
+
+def setup_file_logging(log_file_path: str | Path) -> logging.Handler:
+    file_handler = logging.FileHandler(log_file_path)
+    formatter = logging.Formatter(
+        fmt="{message}",
+        style="{",
+    )
+    file_handler.setFormatter(formatter)
+    file_handler.setLevel(settings.LOG_LEVEL)
+
+    logging.getLogger("birder").addHandler(file_handler)
+    logging.getLogger("birder_clip").addHandler(file_handler)
+
+    return file_handler
+
+
+def save_training_checkpoint(
+    args: argparse.Namespace,
+    network_name: str,
+    epoch: int,
+    net: torch.nn.Module,
+    signature: Any,
+    rgb_stats: Any,
+    optimizer: Optional[torch.optim.Optimizer],
+    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
+    scaler: Optional[torch.amp.grad_scaler.GradScaler],
+    model_base: Optional[torch.nn.Module],
+    *,
+    fsdp_mode: bool = False,
+    fsdp_model_state: Optional[dict[str, Any]] = None,
+    external_config: Optional[dict[str, Any]] = None,
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    if fsdp_mode is True:
+        if fsdp_model_state is not None:
+            model_state = fsdp_model_state
+        else:
+            model_state = fsdp_utils.gather_full_model_state_dict(net)
+
+        optimizer_state = None
+        scheduler_state = None
+        scaler_state = None
+        model_base_state = None
+        if optimizer is not None and scheduler is not None:
+            optimizer_state = fsdp_utils.gather_full_optimizer_state_dict(net, optimizer)
+            scheduler_state = scheduler.state_dict()
+            if scaler is not None:
+                scaler_state = scaler.state_dict()
+            if model_base is not None:
+                model_base_state = model_base.state_dict()
+
+        if birder_training_utils.is_global_primary(args) is True:
+            fs_ops.checkpoint_model_from_state_dicts(
+                network_name,
+                epoch,
+                model_state,
+                net.task,
+                signature,
+                rgb_stats,
+                optimizer_state,
+                scheduler_state,
+                scaler_state,
+                model_base_state,
+                external_config=external_config,
+                **extra_states,
+            )
+
+        if birder_training_utils.is_dist_available_and_initialized() is True:
+            dist.barrier()
+
+    else:
+        if birder_training_utils.is_global_primary(args) is True:
+            fs_ops.checkpoint_model(
+                network_name,
+                epoch,
+                net,
+                signature,
+                rgb_stats,
+                optimizer,
+                scheduler,
+                scaler,
+                model_base,
+                external_config=external_config,
+                **extra_states,
+            )