diff --git a/nemo_curator/models/cosmos_embed1.py b/nemo_curator/models/cosmos_embed1.py index 758c0021f3..86bdcda158 100644 --- a/nemo_curator/models/cosmos_embed1.py +++ b/nemo_curator/models/cosmos_embed1.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib +import importlib.util +from collections.abc import Callable from pathlib import Path from typing import Final, Literal, cast @@ -38,6 +41,191 @@ } +_HEAD_MASK_VECTOR_DIM: Final[int] = 1 +_HEAD_MASK_MATRIX_DIM: Final[int] = 2 +_HEAD_MASK_TARGET_DIM: Final[int] = 5 + + +def _resolve_symbol(name: str) -> object | None: + for module_name in ("transformers.pytorch_utils", "transformers.modeling_utils"): + if importlib.util.find_spec(module_name) is None: + continue + module = importlib.import_module(module_name) + if hasattr(module, name): + return getattr(module, name) + return None + + +def _fallback_apply_chunking_to_forward( + forward_fn: Callable[..., object], chunk_size: int, chunk_dim: int, *input_tensors: torch.Tensor +) -> object: + if len(input_tensors) == 0: + msg = "input_tensors cannot be empty" + raise ValueError(msg) + + if chunk_size <= 0: + return forward_fn(*input_tensors) + + tensor_shape = input_tensors[0].shape[chunk_dim] + if any(t.shape[chunk_dim] != tensor_shape for t in input_tensors): + msg = "All input tensors must have the same shape on chunk_dim" + raise ValueError(msg) + if tensor_shape % chunk_size != 0: + msg = "The dimension to be chunked must be a multiple of chunk_size" + raise ValueError(msg) + + num_chunks = tensor_shape // chunk_size + input_chunks = tuple(t.chunk(num_chunks, dim=chunk_dim) for t in input_tensors) + output_chunks = tuple(forward_fn(*chunk_inputs) for chunk_inputs in zip(*input_chunks, strict=False)) + return torch.cat(output_chunks, dim=chunk_dim) + + +def _fallback_find_pruneable_heads_and_indices( + heads: list[int] | set[int], n_heads: int, head_size: int, already_pruned_heads: set[int] +) -> tuple[set[int], torch.Tensor]: + heads_to_prune = set(heads) - already_pruned_heads + mask = torch.ones(n_heads, head_size) + for head in heads_to_prune: + adjusted_head = head - sum(1 for h in already_pruned_heads if h < head) + mask[adjusted_head] = 0 + mask = mask.view(-1).contiguous().eq(1) + index = torch.arange(mask.numel())[mask].long() + return heads_to_prune, index + + +def _fallback_prune_linear_layer(layer: torch.nn.Linear, index: torch.Tensor, dim: int = 0) -> torch.nn.Linear: + index = index.to(layer.weight.device) + weight = layer.weight.index_select(dim, index).clone().detach() + if layer.bias is not None: + bias = layer.bias.clone().detach() if dim == 1 else layer.bias.index_select(0, index).clone().detach() + new_size = list(layer.weight.size()) + new_size[dim] = index.size(0) + new_layer = torch.nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device) + new_layer.weight.requires_grad = False + new_layer.weight.copy_(weight.contiguous()) + new_layer.weight.requires_grad = True + if layer.bias is not None: + new_layer.bias.requires_grad = False + new_layer.bias.copy_(bias.contiguous()) + new_layer.bias.requires_grad = True + return new_layer + + +def _ensure_all_tied_weights_keys(model: object) -> None: + if hasattr(model, "all_tied_weights_keys"): + return + tied = getattr(model, "_tied_weights_keys", None) + if tied is None: + model.all_tied_weights_keys = {} + return + if isinstance(tied, dict): + model.all_tied_weights_keys = tied + return + try: + model.all_tied_weights_keys = dict.fromkeys(tied) + except TypeError: + model.all_tied_weights_keys = {} + + +def _convert_head_mask_to_5d(self: object, head_mask: torch.Tensor, num_hidden_layers: int) -> torch.Tensor: + if head_mask.dim() == _HEAD_MASK_VECTOR_DIM: + head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) + head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) + elif head_mask.dim() == _HEAD_MASK_MATRIX_DIM: + head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) + if head_mask.dim() != _HEAD_MASK_TARGET_DIM: + msg = f"head_mask.dim != {_HEAD_MASK_TARGET_DIM}, got {head_mask.dim()}" + raise ValueError(msg) + try: + dtype = self.dtype + except Exception: # noqa: BLE001 + dtype = next(self.parameters()).dtype + return head_mask.to(dtype=dtype) + + +def _get_head_mask( + self: object, head_mask: torch.Tensor | None, num_hidden_layers: int, is_attention_chunked: bool = False +) -> torch.Tensor | list[None]: + if head_mask is not None: + head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) + if is_attention_chunked: + head_mask = head_mask.unsqueeze(-1) + return head_mask + return [None] * num_hidden_layers + + +def _patch_modeling_utils_symbols(modeling_utils_module: object) -> None: + if not hasattr(modeling_utils_module, "apply_chunking_to_forward"): + symbol = _resolve_symbol("apply_chunking_to_forward") or _fallback_apply_chunking_to_forward + modeling_utils_module.apply_chunking_to_forward = symbol + if not hasattr(modeling_utils_module, "find_pruneable_heads_and_indices"): + symbol = _resolve_symbol("find_pruneable_heads_and_indices") or _fallback_find_pruneable_heads_and_indices + modeling_utils_module.find_pruneable_heads_and_indices = symbol + if not hasattr(modeling_utils_module, "prune_linear_layer"): + symbol = _resolve_symbol("prune_linear_layer") or _fallback_prune_linear_layer + modeling_utils_module.prune_linear_layer = symbol + + +def _patch_pretrained_model_compat(pre_trained_model_cls: object) -> None: + if not getattr(pre_trained_model_cls, "_nemo_curator_tied_weights_compat_patched", False): + original_mark_tied = getattr(pre_trained_model_cls, "mark_tied_weights_as_initialized", None) + if original_mark_tied is not None: + + def _patched_mark_tied_weights_as_initialized(self: object, *args: object, **kwargs: object) -> object: + _ensure_all_tied_weights_keys(self) + return original_mark_tied(self, *args, **kwargs) + + pre_trained_model_cls.mark_tied_weights_as_initialized = _patched_mark_tied_weights_as_initialized + + original_adjust_tied = getattr(pre_trained_model_cls, "_adjust_tied_keys_with_tied_pointers", None) + if original_adjust_tied is not None: + + def _patched_adjust_tied_keys_with_tied_pointers(self: object, *args: object, **kwargs: object) -> object: + _ensure_all_tied_weights_keys(self) + return original_adjust_tied(self, *args, **kwargs) + + pre_trained_model_cls._adjust_tied_keys_with_tied_pointers = _patched_adjust_tied_keys_with_tied_pointers + pre_trained_model_cls._nemo_curator_tied_weights_compat_patched = True + + if not hasattr(pre_trained_model_cls, "_convert_head_mask_to_5d"): + pre_trained_model_cls._convert_head_mask_to_5d = _convert_head_mask_to_5d + if not hasattr(pre_trained_model_cls, "get_head_mask"): + pre_trained_model_cls.get_head_mask = _get_head_mask + + +def _patch_transformers_compat() -> None: + """Patch moved Transformers symbols expected by older Cosmos remote code.""" + from transformers import modeling_utils + from transformers.modeling_utils import PreTrainedModel + + _patch_modeling_utils_symbols(modeling_utils) + _patch_pretrained_model_compat(PreTrainedModel) + + +def _patch_cosmos_embed1_modeling_vit(*, weights_dir: str, variant: Literal["224p", "336p", "448p"]) -> None: + """Patch Cosmos-Embed1 ViT code to avoid meta-tensor .item() failures.""" + old_line = "dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule" + new_line = 'dpr = torch.linspace(0, drop_path_rate, depth, device="cpu", dtype=torch.float32).tolist()' + candidates = [ + Path(weights_dir) / "modeling_vit.py", + Path.home() + / ".cache/huggingface/modules/transformers_modules" + / f"Cosmos_hyphen_Embed1_hyphen_{variant}" + / "modeling_vit.py", + ] + + for path in candidates: + try: + if not path.exists(): + continue + text = path.read_text(encoding="utf-8") + if old_line not in text: + continue + path.write_text(text.replace(old_line, new_line), encoding="utf-8") + except Exception as e: # noqa: BLE001 + logger.warning(f"Could not patch {path}: {e}") + + class CosmosEmbed1(ModelInterface): """Cosmos-Embed1 embedding model.""" @@ -78,6 +266,8 @@ def setup(self) -> None: This method initializes the model and its configuration for processing video and text data. """ logger.info("Setting up Cosmos-Embed1 model") + _patch_transformers_compat() + _patch_cosmos_embed1_modeling_vit(weights_dir=self._weights_dir, variant=self.variant) if not Path(self._weights_dir).exists(): exception = f"Weights directory {self._weights_dir} not found!" raise FileNotFoundError(exception) @@ -201,7 +391,9 @@ def download_weights_on_node(cls, model_dir: str, variant: Literal["224p", "336p logger.info(f"CosmosEmbed1 {variant} weights downloaded to: {model_dir_path}") @classmethod - def download_processor_config_on_node(cls, model_dir: str, variant: Literal["224p", "336p", "448p"] = "336p") -> None: + def download_processor_config_on_node( + cls, model_dir: str, variant: Literal["224p", "336p", "448p"] = "336p" + ) -> None: """Download the processor config for the CosmosEmbed1 model on the node.""" model_dir_path = Path(model_dir) / _COSMOS_EMBED1_VARIANTS_INFO[variant] model_dir_path.mkdir(parents=True, exist_ok=True)