Add DoL 1.2 model card and inference runtime

431ffa3 verified 3 months ago

8.49 kB

# -*- coding: utf-8 -*-

from __future__ import annotations

from typing import Mapping

import torch
import torch.nn as nn


class MLPVisionEncoder(nn.Module):
    """Lightweight MLP+LN patch encoder.
    Kept as a fallback for smoke tests and small-data experiments; production
    VLM training should use ``Model.omvt.OMVTVisionTower`` instead.
    """

    def __init__(self, cfg, patch_pixels: int = 14 * 14 * 3):
        super().__init__()

        if patch_pixels <= 0:
            raise ValueError("patch_pixels must be positive")

        self.cfg = cfg
        self.patch_pixels = patch_pixels
        self.d_model = cfg.d_model

        self.patch_embed = nn.Linear(patch_pixels, cfg.d_model)

        self.encoder = nn.Sequential(
            nn.LayerNorm(cfg.d_model),
            nn.Linear(cfg.d_model, cfg.d_model * 2),
            nn.GELU(),
            nn.Linear(cfg.d_model * 2, cfg.d_model),
            nn.LayerNorm(cfg.d_model),
        )

        self.to_llm = nn.Linear(cfg.d_model, cfg.d_model)

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        if pixel_values.ndim not in {2, 3}:
            raise ValueError("pixel_values must have shape [N, P] or [B, N, P]")

        if pixel_values.shape[-1] != self.patch_pixels:
            raise ValueError(
                f"expected patch_pixels={self.patch_pixels}, got {pixel_values.shape[-1]}"
            )

        x = self.patch_embed(pixel_values)
        x = x + self.encoder(x)
        return self.to_llm(x)


# Back-compat alias; legacy code expecting VisionEncoder still works.
VisionEncoder = MLPVisionEncoder


def inject_visual_features(
    inputs_embeds: torch.Tensor,
    input_ids: torch.Tensor,
    visual_features: torch.Tensor,
    image_patch_id: int,
) -> torch.Tensor:
    if inputs_embeds.ndim != 3:
        raise ValueError("inputs_embeds must have shape [B, L, D]")

    if input_ids.ndim != 2:
        raise ValueError("input_ids must have shape [B, L]")

    if input_ids.shape != inputs_embeds.shape[:2]:
        raise ValueError("input_ids shape must match inputs_embeds[:2]")

    if visual_features.ndim not in {2, 3}:
        raise ValueError("visual_features must have shape [N, D] or [B, N, D]")

    bsz, _seq_len, dim = inputs_embeds.shape
    mask = input_ids == image_patch_id
    n_patches = int(mask.sum().item())

    if n_patches == 0:
        return inputs_embeds

    out = inputs_embeds.clone()

    if visual_features.ndim == 2:
        if bsz != 1:
            raise ValueError(
                "unbatched visual_features are only supported when batch size is 1"
            )

        if visual_features.shape != (n_patches, dim):
            raise ValueError(
                f"expected visual_features {(n_patches, dim)}, got {tuple(visual_features.shape)}"
            )

        out[mask] = visual_features.to(device=out.device, dtype=out.dtype)
        return out

    if visual_features.shape[0] != bsz or visual_features.shape[-1] != dim:
        raise ValueError(
            f"expected visual_features [B, N, {dim}], got {tuple(visual_features.shape)}"
        )

    counts = mask.sum(dim=1).tolist()

    for batch_idx, count in enumerate(counts):
        if visual_features.shape[1] != count:
            raise ValueError(
                f"batch {batch_idx}: visual feature count must equal image_patch "
                f"count (got {visual_features.shape[1]} vs {count})"
            )

        if count > 0:
            out[batch_idx, mask[batch_idx]] = visual_features[
                batch_idx,
                :count,
            ].to(device=out.device, dtype=out.dtype)

    return out


class VisionInjector(nn.Module):
    """Dispatcher between MLP fallback and OMVT vision tower.
    When ``pixel_values`` is a ``Tensor``, the lightweight :class:`MLPVisionEncoder`
    is used (legacy / smoke path). When it is a ``Mapping`` (the multi-scale
    batch produced by :class:`Model.omvt.MultiScalePatcher`), the OMVT tower
    + Perceiver compressor are dispatched instead.
    The OMVT tower is constructed lazily on the first dict input so smoke runs
    that never see vision input pay no parameter cost.
    """

    def __init__(
        self,
        cfg,
        patch_pixels: int = 14 * 14 * 3,
        omvt_cfg=None,
    ):
        super().__init__()

        self.cfg = cfg
        self.patch_pixels = patch_pixels
        self.encoder = MLPVisionEncoder(cfg, patch_pixels)

        self._omvt_cfg = omvt_cfg
        self.omvt = None  # type: ignore[assignment]

    def _ensure_omvt(self) -> None:
        if self.omvt is not None:
            return
        if self._omvt_cfg is None:
            from Model.config import OMVTConfig

            self._omvt_cfg = OMVTConfig()

        from Model.omvt import OMVTInjector

        injector = OMVTInjector(self.cfg, self._omvt_cfg)

        # Lazy submodules must be migrated to the parent's current device
        # and dtype: assigning after `RDTForCausalLM.to('cuda')` does not
        # auto-migrate, so the first GPU batch would crash with a device
        # mismatch. Infer the target from an existing parameter; fall back
        # to MLP encoder's first parameter, then CPU/float32 if the module
        # is genuinely parameterless (shouldn't happen in practice).
        target_param = next(self.parameters(), None)
        if target_param is not None:
            injector = injector.to(device=target_param.device, dtype=target_param.dtype)

        self.omvt = injector

    def forward(
        self,
        inputs_embeds: torch.Tensor,
        input_ids: torch.Tensor,
        pixel_values: torch.Tensor | Mapping[str, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        if pixel_values is None:
            return inputs_embeds

        if isinstance(pixel_values, Mapping):
            self._ensure_omvt()
            visual_features = self.omvt(pixel_values)
            return inject_visual_features(
                inputs_embeds=inputs_embeds,
                input_ids=input_ids,
                visual_features=visual_features,
                image_patch_id=self.cfg.image_patch_id,
            )

        visual_features = self.encoder(pixel_values)
        return inject_visual_features(
            inputs_embeds=inputs_embeds,
            input_ids=input_ids,
            visual_features=visual_features,
            image_patch_id=self.cfg.image_patch_id,
        )


def _check() -> None:
    from Model.config import tiny_config

    torch.manual_seed(0)

    cfg = tiny_config()
    patch_pixels = 14 * 14 * 3

    input_ids = torch.tensor(
        [
            [
                cfg.bos_id,
                300,
                301,
                cfg.image_start_id,
                cfg.image_patch_id,
                cfg.image_patch_id,
                cfg.image_patch_id,
                cfg.image_patch_id,
                cfg.image_end_id,
                302,
                303,
                cfg.eos_id,
            ]
        ]
    )

    inputs_embeds = torch.randn(1, input_ids.shape[1], cfg.d_model)
    pixel_values = torch.randn(4, patch_pixels)

    injector = VisionInjector(cfg, patch_pixels)
    out = injector(inputs_embeds, input_ids, pixel_values)

    mask = input_ids == cfg.image_patch_id
    changed = (out != inputs_embeds).any(dim=-1)

    print("VisionInjector")
    print(f"  shape: {tuple(inputs_embeds.shape)} -> {tuple(out.shape)}")
    print(f"  patch_count: {int(mask.sum().item())}")
    print(f"  changed: {changed[0].nonzero().squeeze(-1).tolist()}")
    print(f"  expected: {mask[0].nonzero().squeeze(-1).tolist()}")
    print(f"  exact_match: {torch.equal(changed, mask)}")

    out_text = injector(inputs_embeds, input_ids, pixel_values=None)
    print(f"  text_only_equal: {torch.equal(out_text, inputs_embeds)}")

    pv = torch.randn(4, patch_pixels, requires_grad=True)
    out = injector(inputs_embeds, input_ids, pv)
    out.sum().backward()

    print(f"  grad_norm: {pv.grad.norm().item():.6f}")

    batched_ids = torch.tensor(
        [
            [cfg.bos_id, cfg.image_patch_id, cfg.image_patch_id, cfg.eos_id],
            [cfg.bos_id, 301, cfg.image_patch_id, cfg.eos_id],
        ]
    )
    batched_embeds = torch.randn(2, 4, cfg.d_model)
    batched_pixels = torch.randn(2, 2, patch_pixels)

    batched_out = injector(batched_embeds, batched_ids, batched_pixels)
    print(f"  batched_shape: {tuple(batched_out.shape)}")


if __name__ == "__main__":
    _check()

	# -- coding: utf-8 --

	from __future__ import annotations

	from typing import Mapping

	import torch
	import torch.nn as nn


	class MLPVisionEncoder(nn.Module):
	"""Lightweight MLP+LN patch encoder.

	Kept as a fallback for smoke tests and small-data experiments; production
	VLM training should use ``Model.omvt.OMVTVisionTower`` instead.
	"""

	def __init__(self, cfg, patch_pixels: int = 14 * 14 * 3):
	super().__init__()

	if patch_pixels <= 0:
	raise ValueError("patch_pixels must be positive")

	self.cfg = cfg
	self.patch_pixels = patch_pixels
	self.d_model = cfg.d_model

	self.patch_embed = nn.Linear(patch_pixels, cfg.d_model)

	self.encoder = nn.Sequential(
	nn.LayerNorm(cfg.d_model),
	nn.Linear(cfg.d_model, cfg.d_model * 2),
	nn.GELU(),
	nn.Linear(cfg.d_model * 2, cfg.d_model),
	nn.LayerNorm(cfg.d_model),
	)

	self.to_llm = nn.Linear(cfg.d_model, cfg.d_model)

	def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
	if pixel_values.ndim not in {2, 3}:
	raise ValueError("pixel_values must have shape [N, P] or [B, N, P]")

	if pixel_values.shape[-1] != self.patch_pixels:
	raise ValueError(
	f"expected patch_pixels={self.patch_pixels}, got {pixel_values.shape[-1]}"
	)

	x = self.patch_embed(pixel_values)
	x = x + self.encoder(x)
	return self.to_llm(x)


	# Back-compat alias; legacy code expecting VisionEncoder still works.
	VisionEncoder = MLPVisionEncoder


	def inject_visual_features(
	inputs_embeds: torch.Tensor,
	input_ids: torch.Tensor,
	visual_features: torch.Tensor,
	image_patch_id: int,
	) -> torch.Tensor:
	if inputs_embeds.ndim != 3:
	raise ValueError("inputs_embeds must have shape [B, L, D]")

	if input_ids.ndim != 2:
	raise ValueError("input_ids must have shape [B, L]")

	if input_ids.shape != inputs_embeds.shape[:2]:
	raise ValueError("input_ids shape must match inputs_embeds[:2]")

	if visual_features.ndim not in {2, 3}:
	raise ValueError("visual_features must have shape [N, D] or [B, N, D]")

	bsz, _seq_len, dim = inputs_embeds.shape
	mask = input_ids == image_patch_id
	n_patches = int(mask.sum().item())

	if n_patches == 0:
	return inputs_embeds

	out = inputs_embeds.clone()

	if visual_features.ndim == 2:
	if bsz != 1:
	raise ValueError(
	"unbatched visual_features are only supported when batch size is 1"
	)

	if visual_features.shape != (n_patches, dim):
	raise ValueError(
	f"expected visual_features {(n_patches, dim)}, got {tuple(visual_features.shape)}"
	)

	out[mask] = visual_features.to(device=out.device, dtype=out.dtype)
	return out

	if visual_features.shape[0] != bsz or visual_features.shape[-1] != dim:
	raise ValueError(
	f"expected visual_features [B, N, {dim}], got {tuple(visual_features.shape)}"
	)

	counts = mask.sum(dim=1).tolist()

	for batch_idx, count in enumerate(counts):
	if visual_features.shape[1] != count:
	raise ValueError(
	f"batch {batch_idx}: visual feature count must equal image_patch "
	f"count (got {visual_features.shape[1]} vs {count})"
	)

	if count > 0:
	out[batch_idx, mask[batch_idx]] = visual_features[
	batch_idx,
	:count,
	].to(device=out.device, dtype=out.dtype)

	return out


	class VisionInjector(nn.Module):
	"""Dispatcher between MLP fallback and OMVT vision tower.

	When ``pixel_values`` is a ``Tensor``, the lightweight :class:`MLPVisionEncoder`
	is used (legacy / smoke path). When it is a ``Mapping`` (the multi-scale
	batch produced by :class:`Model.omvt.MultiScalePatcher`), the OMVT tower
	+ Perceiver compressor are dispatched instead.

	The OMVT tower is constructed lazily on the first dict input so smoke runs
	that never see vision input pay no parameter cost.
	"""

	def __init__(
	self,
	cfg,
	patch_pixels: int = 14 * 14 * 3,
	omvt_cfg=None,
	):
	super().__init__()

	self.cfg = cfg
	self.patch_pixels = patch_pixels
	self.encoder = MLPVisionEncoder(cfg, patch_pixels)

	self._omvt_cfg = omvt_cfg
	self.omvt = None # type: ignore[assignment]

	def _ensure_omvt(self) -> None:
	if self.omvt is not None:
	return
	if self._omvt_cfg is None:
	from Model.config import OMVTConfig

	self._omvt_cfg = OMVTConfig()

	from Model.omvt import OMVTInjector

	injector = OMVTInjector(self.cfg, self._omvt_cfg)

	# Lazy submodules must be migrated to the parent's current device
	# and dtype: assigning after `RDTForCausalLM.to('cuda')` does not
	# auto-migrate, so the first GPU batch would crash with a device
	# mismatch. Infer the target from an existing parameter; fall back
	# to MLP encoder's first parameter, then CPU/float32 if the module
	# is genuinely parameterless (shouldn't happen in practice).
	target_param = next(self.parameters(), None)
	if target_param is not None:
	injector = injector.to(device=target_param.device, dtype=target_param.dtype)

	self.omvt = injector

	def forward(
	self,
	inputs_embeds: torch.Tensor,
	input_ids: torch.Tensor,
	pixel_values: torch.Tensor \| Mapping[str, torch.Tensor] \| None = None,
	) -> torch.Tensor:
	if pixel_values is None:
	return inputs_embeds

	if isinstance(pixel_values, Mapping):
	self._ensure_omvt()
	visual_features = self.omvt(pixel_values)
	return inject_visual_features(
	inputs_embeds=inputs_embeds,
	input_ids=input_ids,
	visual_features=visual_features,
	image_patch_id=self.cfg.image_patch_id,
	)

	visual_features = self.encoder(pixel_values)
	return inject_visual_features(
	inputs_embeds=inputs_embeds,
	input_ids=input_ids,
	visual_features=visual_features,
	image_patch_id=self.cfg.image_patch_id,
	)


	def _check() -> None:
	from Model.config import tiny_config

	torch.manual_seed(0)

	cfg = tiny_config()
	patch_pixels = 14 * 14 * 3

	input_ids = torch.tensor(
	[
	[
	cfg.bos_id,
	300,
	301,
	cfg.image_start_id,
	cfg.image_patch_id,
	cfg.image_patch_id,
	cfg.image_patch_id,
	cfg.image_patch_id,
	cfg.image_end_id,
	302,
	303,
	cfg.eos_id,
	]
	]
	)

	inputs_embeds = torch.randn(1, input_ids.shape[1], cfg.d_model)
	pixel_values = torch.randn(4, patch_pixels)

	injector = VisionInjector(cfg, patch_pixels)
	out = injector(inputs_embeds, input_ids, pixel_values)

	mask = input_ids == cfg.image_patch_id
	changed = (out != inputs_embeds).any(dim=-1)

	print("VisionInjector")
	print(f" shape: {tuple(inputs_embeds.shape)} -> {tuple(out.shape)}")
	print(f" patch_count: {int(mask.sum().item())}")
	print(f" changed: {changed[0].nonzero().squeeze(-1).tolist()}")
	print(f" expected: {mask[0].nonzero().squeeze(-1).tolist()}")
	print(f" exact_match: {torch.equal(changed, mask)}")

	out_text = injector(inputs_embeds, input_ids, pixel_values=None)
	print(f" text_only_equal: {torch.equal(out_text, inputs_embeds)}")

	pv = torch.randn(4, patch_pixels, requires_grad=True)
	out = injector(inputs_embeds, input_ids, pv)
	out.sum().backward()

	print(f" grad_norm: {pv.grad.norm().item():.6f}")

	batched_ids = torch.tensor(
	[
	[cfg.bos_id, cfg.image_patch_id, cfg.image_patch_id, cfg.eos_id],
	[cfg.bos_id, 301, cfg.image_patch_id, cfg.eos_id],
	]
	)
	batched_embeds = torch.randn(2, 4, cfg.d_model)
	batched_pixels = torch.randn(2, 2, patch_pixels)

	batched_out = injector(batched_embeds, batched_ids, batched_pixels)
	print(f" batched_shape: {tuple(batched_out.shape)}")


	if __name__ == "__main__":
	_check()