backbone_adapters.py

"""
Backbone adapters for diffusion retriever training.

Each adapter encapsulates ALL model-specific behavior in one place:
  - How to load the backbone (AutoModel vs AutoModelForCausalLM)
  - PEFT/LoRA configuration (target modules, task type)
  - Attention mask format (2D vs 4D)
  - Hidden state extraction (forward hook on output projection)
  - Mask token ID (verified from HuggingFace tokenizer configs)
  - Gradient checkpointing support

The TrainableDiffusionRetriever delegates to an adapter and has ZERO
model-specific branches.  Adding a new model = adding one adapter class.
"""

from abc import ABC, abstractmethod
from typing import Dict, Optional, Tuple
import logging

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Base adapter
# ---------------------------------------------------------------------------

class BackboneAdapter(ABC):
    """Abstract interface for model-specific backbone behavior."""

    model_type: str          # e.g. 'dream', 'llada1', 'llada2'
    mask_token_id: int       # verified from HuggingFace model cards
    hub_model_name: str      # HuggingFace model ID for fallback loading

    def __init__(self):
        self.flash_attn: bool = False   # set by load_backbone

    # -- loading --------------------------------------------------------

    def load_backbone(self, source: str, device_map=None) -> nn.Module:
        """Load backbone from HuggingFace model name or local directory.

        Tries flash attention variants in order, falls back to eager.
        Sets ``self.flash_attn`` as a side effect.
        """
        common_kw = dict(trust_remote_code=True, torch_dtype=torch.bfloat16)
        if device_map is not None:
            common_kw['device_map'] = device_map

        for attn_impl in self._flash_attn_impls:
            try:
                bb = self._auto_class().from_pretrained(
                    source, attn_implementation=attn_impl, **common_kw)
                self.flash_attn = True
                logger.info(f"{self.model_type}: {attn_impl} enabled")
                return bb
            except (ValueError, ImportError):
                pass

        self.flash_attn = False
        return self._auto_class().from_pretrained(source, **common_kw)

    @staticmethod
    @abstractmethod
    def _auto_class():
        """Return the AutoModel class to use (AutoModel or AutoModelForCausalLM)."""

    _flash_attn_impls: Tuple[str, ...] = ('flash_attention_2',)

    # -- PEFT / LoRA ----------------------------------------------------

    @abstractmethod
    def get_lora_config(self, lora_rank: int, lora_alpha: int,
                        lora_dropout: float = 0.0):
        """Return a ``peft.LoraConfig`` appropriate for this backbone."""

    # -- attention mask --------------------------------------------------

    @abstractmethod
    def needs_4d_mask(self) -> bool:
        """Whether the backbone expects a 4D ``[B,1,L,L]`` attention mask.

        If False, the backbone handles bidirectional attention internally
        and expects a standard 2D ``[B,L]`` padding mask.
        """

    # -- hidden state extraction -----------------------------------------

    def register_hidden_hook(self, backbone: nn.Module,
                             ref_dict: Dict[str, torch.Tensor]) -> bool:
        """Register a forward hook on the output projection to capture the
        last hidden state without ``output_hidden_states=True``.

        Returns True if a hook was registered, False otherwise (in which
        case the caller should fall back to ``output_hidden_states``).
        """
        return False  # default: no hook, use output_hidden_states

    # -- gradient checkpointing ------------------------------------------

    def enable_gradient_checkpointing(self, backbone: nn.Module, **kwargs):
        """Enable gradient checkpointing.  Override for models that don't
        support it."""
        backbone.gradient_checkpointing_enable(**kwargs)
        logger.info("Gradient checkpointing enabled")


# ---------------------------------------------------------------------------
# Hook helpers (shared across adapters)
# ---------------------------------------------------------------------------

def _is_linear(mod: nn.Module) -> bool:
    """Check if a module is a Linear layer (plain or LoRA-wrapped)."""
    if isinstance(mod, nn.Linear):
        return True
    # PEFT LoRA wraps nn.Linear in peft.tuners.lora.layer.Linear which is
    # NOT a subclass of nn.Linear, but has a base_layer that is.
    if hasattr(mod, 'base_layer') and isinstance(mod.base_layer, nn.Linear):
        return True
    return False


def _hook_on_module(backbone: nn.Module, ref_dict: Dict,
                    target_name: str, skip_if_contains: Optional[str] = None,
                    adapter_name: str = '') -> bool:
    """Register a forward hook on the first Linear (or LoRA-wrapped Linear)
    whose leaf name matches *target_name*.  Optionally skip modules whose
    full path contains *skip_if_contains* (e.g. 'blocks' to skip per-layer
    ff_out).
    """
    for name, mod in backbone.named_modules():
        leaf = name.split('.')[-1]
        if leaf == target_name and _is_linear(mod):
            if skip_if_contains and skip_if_contains in name:
                continue
            mod.register_forward_hook(
                lambda m, inp, out, r=ref_dict: r.update({'h': inp[0]})
            )
            logger.info(f"{adapter_name}: hook on '{name}'")
            return True
    return False


# ---------------------------------------------------------------------------
# Dream
# ---------------------------------------------------------------------------

class DreamAdapter(BackboneAdapter):
    model_type = 'dream'
    mask_token_id = 151666          # <|mask|>  — Dream-org/Dream-v0-Instruct-7B
    hub_model_name = 'Dream-org/Dream-v0-Instruct-7B'

    @staticmethod
    def _auto_class():
        from transformers import AutoModel
        return AutoModel

    def get_lora_config(self, lora_rank, lora_alpha, lora_dropout=0.0):
        from peft import LoraConfig, TaskType
        # Dream is loaded via AutoModel (not AutoModelForCausalLM), so
        # FEATURE_EXTRACTION avoids PeftModelForCausalLM which would
        # require prepare_inputs_for_generation (Dream lacks this).
        return LoraConfig(
            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                            "gate_proj", "up_proj", "down_proj"],
            task_type=TaskType.FEATURE_EXTRACTION,
            bias="none",
        )

    def needs_4d_mask(self) -> bool:
        # Standard HF Qwen2 attention — needs 4D to enforce bidirectional.
        return True

    def register_hidden_hook(self, backbone, ref_dict):
        return _hook_on_module(backbone, ref_dict, 'lm_head',
                               adapter_name='dream')


# ---------------------------------------------------------------------------
# LLaDA v1 (GSAI-ML/LLaDA-8B-Instruct)
# ---------------------------------------------------------------------------

class LLaDA1Adapter(BackboneAdapter):
    model_type = 'llada1'
    mask_token_id = 126336          # <|mdm_mask|>  — GSAI-ML/LLaDA-8B-Instruct
    hub_model_name = 'GSAI-ML/LLaDA-8B-Instruct'
    _flash_attn_impls = ('flash_attention_3', 'flash_attention_2')

    @staticmethod
    def _auto_class():
        from transformers import AutoModelForCausalLM
        return AutoModelForCausalLM

    def get_lora_config(self, lora_rank, lora_alpha, lora_dropout=0.0):
        from peft import LoraConfig, TaskType
        # LLaDA1 custom arch uses different names than standard LLaMA:
        #   attn_out (not o_proj), ff_proj (not gate_proj), ff_out (not down_proj)
        return LoraConfig(
            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
            target_modules=["q_proj", "k_proj", "v_proj", "attn_out",
                            "ff_proj", "up_proj", "ff_out"],
            task_type=TaskType.CAUSAL_LM,
            bias="none",
        )

    def needs_4d_mask(self) -> bool:
        # Custom code handles bidirectional attention internally — always 2D.
        return False

    def register_hidden_hook(self, backbone, ref_dict):
        # LLaDA1 has ff_out per-block (blocks.X.ff_out = FFN output) AND at
        # the model level (transformer.ff_out = output projection to vocab).
        # We need the model-level one; skip per-block ones.
        return _hook_on_module(backbone, ref_dict, 'ff_out',
                               skip_if_contains='blocks',
                               adapter_name='llada1')

    def enable_gradient_checkpointing(self, backbone, **kwargs):
        # LLaDA1's LLaDAModelLM doesn't support HF gradient_checkpointing_enable.
        # Manually wrap each transformer block with torch checkpoint.
        from torch.utils.checkpoint import checkpoint as ckpt_fn

        # Find the blocks ModuleList through the PEFT wrapper
        blocks = None
        for name, mod in backbone.named_modules():
            if name.endswith('.blocks') and isinstance(mod, nn.ModuleList):
                blocks = mod
                break

        if blocks is None:
            logger.warning("LLaDA1: couldn't find transformer blocks — "
                           "skipping gradient checkpointing")
            return

        for block in blocks:
            orig_forward = block.forward

            def _make_ckpt(fwd):
                def _ckpt_forward(*args, **kwargs):
                    if not torch.is_grad_enabled():
                        return fwd(*args, **kwargs)
                    return ckpt_fn(fwd, *args, use_reentrant=False, **kwargs)
                return _ckpt_forward

            block.forward = _make_ckpt(orig_forward)

        logger.info(f"LLaDA1: manual gradient checkpointing on {len(blocks)} blocks")


# ---------------------------------------------------------------------------
# LLaDA v1.5 (GSAI-ML/LLaDA-1.5)  — same architecture as v1
# ---------------------------------------------------------------------------

class LLaDA15Adapter(LLaDA1Adapter):
    model_type = 'llada15'
    mask_token_id = 126336          # <|mdm_mask|>  — same tokenizer as v1
    hub_model_name = 'GSAI-ML/LLaDA-1.5'

    def register_hidden_hook(self, backbone, ref_dict):
        return _hook_on_module(backbone, ref_dict, 'ff_out',
                               skip_if_contains='blocks',
                               adapter_name='llada15')



# ---------------------------------------------------------------------------
# LLaDA v2 (inclusionAI/LLaDA2.0-mini)
# ---------------------------------------------------------------------------

class LLaDA2Adapter(BackboneAdapter):
    model_type = 'llada2'
    mask_token_id = 156895          # <|mask|>  — inclusionAI/LLaDA2.0-mini
    hub_model_name = 'inclusionAI/LLaDA2.0-mini'
    _flash_attn_impls = ('flash_attention_3', 'flash_attention_2')

    @staticmethod
    def _auto_class():
        from transformers import AutoModelForCausalLM
        return AutoModelForCausalLM

    def get_lora_config(self, lora_rank, lora_alpha, lora_dropout=0.0):
        from peft import LoraConfig, TaskType
        # LLaDA2 uses fused QKV ("query_key_value") and "dense" for attn
        # output.  Skip MoE expert FFN layers to avoid multiplying params.
        return LoraConfig(
            r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
            target_modules=[
                "query_key_value", "dense",
                "mlp.shared_experts.gate_proj",
                "mlp.shared_experts.up_proj",
                "mlp.shared_experts.down_proj",
            ],
            task_type=TaskType.CAUSAL_LM,
            bias="none",
        )

    def needs_4d_mask(self) -> bool:
        # Standard HF causal model — needs 4D to override causal attention,
        # unless flash attention handles masking itself.
        return not self.flash_attn

    def register_hidden_hook(self, backbone, ref_dict):
        return _hook_on_module(backbone, ref_dict, 'lm_head',
                               adapter_name='llada2')


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

ADAPTER_REGISTRY: Dict[str, type] = {
    'dream':  DreamAdapter,
    'llada1': LLaDA1Adapter,
    'llada15': LLaDA15Adapter,
    'llada2': LLaDA2Adapter,
}


def get_adapter(model_type: str) -> BackboneAdapter:
    """Create a BackboneAdapter for the given model_type."""
    cls = ADAPTER_REGISTRY.get(model_type)
    if cls is None:
        raise ValueError(
            f"Unknown model_type: {model_type!r}. "
            f"Available: {sorted(ADAPTER_REGISTRY.keys())}")
    return cls()