Spaces:

nvidia
/

DVSM

Running on Zero

Initial commit

cebb76b 5 months ago

9.41 kB

import torch
import torch.nn.functional as F
import torch.utils.checkpoint

# flash-attn is a hard dependency in the training / NVIDIA-container environment,
# but it is unavailable (or ABI/arch-incompatible) on some inference targets such
# as HF Spaces ZeroGPU (torch 2.8+, Blackwell sm_120). Soft-import it and fall
# back to torch SDPA, which is numerically equivalent for our usage (non-causal,
# no dropout, default scale = 1/sqrt(head_dim)).
import os

try:
    import flash_attn
    _HAS_FLASH = True
except Exception:  # ImportError or runtime/arch incompatibility on import
    flash_attn = None
    _HAS_FLASH = False

# Allow forcing the SDPA path even when flash_attn is importable, so the
# ZeroGPU code path can be validated on a flash-attn-equipped GPU node.
if os.environ.get("DVSM_DISABLE_FLASH") == "1":
    _HAS_FLASH = False


def _attn_func(q, k, v):
    """Drop-in replacement for ``flash_attn.flash_attn_func``.
    flash_attn layout: q/k/v are ``[B, L, num_head, head_dim]``, non-causal,
    softmax_scale = 1/sqrt(head_dim). torch SDPA expects ``[B, num_head, L, head_dim]``
    and uses the same default scale + non-causal semantics, so we just transpose
    in and out. Output dtype matches the inputs (bf16/fp16), same as flash_attn.
    """
    if _HAS_FLASH:
        return flash_attn.flash_attn_func(q, k, v)
    qh = q.transpose(1, 2)  # [B, H, Lq, D]
    kh = k.transpose(1, 2)  # [B, H, Lk, D]
    vh = v.transpose(1, 2)  # [B, H, Lk, D]
    out = F.scaled_dot_product_attention(qh, kh, vh)  # default scale, non-causal
    return out.transpose(1, 2).contiguous()  # back to [B, Lq, H, D]


def checkpoint_wrapper(module):
    # Wrap model forward function to automatically use activation checkpointing
    ori_forward = module.forward
    def new_forward(*args, **kwargs):
        if module.training:
            output = torch.utils.checkpoint.checkpoint(
                ori_forward, *args, **kwargs, use_reentrant=False)
        else:
            output = ori_forward(*args, **kwargs)
        return output
    return new_forward


class Contiguous(torch.nn.Module):
    def forward(self, x):
        return x.contiguous()


class MLP(torch.nn.Module):
    def __init__(self, dim, expand_ratio=4, bias=False, activation=torch.nn.GELU):
        """
        Args:
            dim: Input dimension
            expand_ratio: Multiplier for hidden dimension
            bias: Whether to use bias in linear layers
            activation: Activation function
        """
        super().__init__()
        hidden_dim = round(dim * expand_ratio)
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(dim, hidden_dim, bias=bias),
            activation(),
            torch.nn.Linear(hidden_dim, dim, bias=bias))

    def forward(self, x):
        return self.mlp(x)


class SwiGluMLP(torch.nn.Module):
    def __init__(self, dim, expand_ratio=4, bias=False):
        """
        Args:
            dim: Input dimension
            expand_ratio: Multiplier for hidden dimension
            bias: Whether to use bias in linear layers
            activation: Activation function
        """
        super().__init__()
        hidden_dim = round(dim * expand_ratio)
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(dim, hidden_dim, bias=bias),
            torch.nn.Linear(dim, hidden_dim, bias=bias),
            torch.nn.Linear(hidden_dim, dim, bias=bias))

    def forward(self, x):
        gate = torch.nn.functional.silu(self.mlp[0](x), inplace=True)
        return self.mlp[2](gate * self.mlp[1](x))


class IntraViewAttention(torch.nn.Module):
    def __init__(self,
        dim,
        head_dim,
        bias=False,
        use_qk_norm=True,
        use_gating=False,
        use_sandwich=False,
        out_dim=None,
    ):
        """
        dim: Input dimension
        head_dim: Dimension of each attention head
        bias: Whether to use bias in QKV and output projection
        use_qk_norm: Whether to use Q-K normalization
        use_gating: Whether to apply gating before output ffn
        """
        super().__init__()
        assert dim % head_dim == 0, f"dim={dim} should be divisible by head_dim={head_dim}"
        num_head = dim // head_dim

        self.dim = dim
        self.out_dim = dim if out_dim is None else out_dim
        self.head_dim = head_dim
        self.num_head = num_head
        self.use_qk_norm = use_qk_norm
        self.use_gating = use_gating
        self.use_sandwich = use_sandwich
        self.attn_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

        self.to_qkv = torch.nn.Linear(dim, 3 * dim, bias=bias)
        self.fc = torch.nn.Linear(dim, self.out_dim, bias=bias)

        # Optional Q-K normalization
        if self.use_qk_norm:
            self.q_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
            self.k_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
        else:
            self.q_norm = torch.nn.Identity()
            self.k_norm = torch.nn.Identity()

        # Optional gating
        if self.use_gating:
            self.gate = torch.nn.Sequential(
                torch.nn.Linear(dim, dim, bias=False),
                torch.nn.Sigmoid())
        # Optional normalization
        if self.use_sandwich:
            self.sandwich = torch.nn.LayerNorm(dim, bias=False)

    def forward(self, x):
        """
        x: [B, L, D]
        """
        # Predict qkv packed as (batch, seq_len, 3, num_head, head_dim)
        qkv = self.to_qkv(x).unflatten(2, (3, self.num_head, self.head_dim))

        # Unpack (q, k, v) each as (batch, seq_len, num_head, head_dim)
        q, k, v = qkv.moveaxis(2,0).contiguous()

        # Apply qk normalization if enabled
        q = self.q_norm(q)
        k = self.k_norm(k)

        # Attention
        with torch.autocast(enabled=False, device_type="cuda"):
            q = q.to(self.attn_dtype)
            k = k.to(self.attn_dtype)
            v = v.to(self.attn_dtype)

        out = _attn_func(q, k, v)

        # Output layer to yield (batch, seq_len, dim)
        out = out.flatten(2,3)
        if self.use_gating:
            out = out * self.gate(x)
        out = self.fc(out)
        if self.use_sandwich:
            out = self.sandwich(out)

        return out


class CrossViewAttention(torch.nn.Module):
    def __init__(self,
        dim,
        head_dim,
        bias=False,
        use_qk_norm=True,
        use_gating=False,
        use_sandwich=False,
        out_dim=None,
    ):
        """
        Args:
            dim: Input dimension
            head_dim: Dimension of each attention head
            bias: Whether to use bias in QKV and output projection
            use_qk_norm: Whether to use Q-K normalization
            use_gating: Whether to apply gating before output ffn
        """
        super().__init__()
        assert dim % head_dim == 0, f"dim={dim} should be divisible by head_dim={head_dim}"
        num_head = dim // head_dim

        self.dim = dim
        self.out_dim = dim if out_dim is None else out_dim
        self.head_dim = head_dim
        self.num_head = num_head
        self.use_qk_norm = use_qk_norm
        self.use_gating = use_gating
        self.use_sandwich = use_sandwich
        self.attn_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

        self.to_q = torch.nn.Linear(dim, dim, bias=bias)
        self.to_kv = torch.nn.Linear(dim, 2 * dim, bias=bias)
        self.fc = torch.nn.Linear(dim, self.out_dim, bias=bias)

        # Optional Q-K normalization
        if self.use_qk_norm:
            self.q_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
            self.k_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
        else:
            self.q_norm = torch.nn.Identity()
            self.k_norm = torch.nn.Identity()

        # Optional gating
        if self.use_gating:
            self.gate = torch.nn.Sequential(
                torch.nn.Linear(dim, dim, bias=False),
                torch.nn.Sigmoid())
        # Optional normalization
        if self.use_sandwich:
            self.sandwich = torch.nn.LayerNorm(dim, bias=False)

    def forward(self, x, kv=None):
        """
        x: [B, L, D]
        kv: ([B, M, D], [B, M, D]) If provided, x query on kv instead of doing self attention
        """
        q = self.to_q(x)

        # To shape of (batch, seq_len, num_head, head_dim)
        q = q.unflatten(2, (self.num_head, self.head_dim))
        q = self.q_norm(q)
        with torch.autocast(enabled=False, device_type="cuda"):
            q = q.to(self.attn_dtype)

        # Prepare (k, v)
        if kv is None:
            kv = self.to_kv(x)
            kv = kv.unflatten(2, (2, self.num_head, self.head_dim))
            k, v = kv.moveaxis(2, 0).contiguous()
            k = self.k_norm(k)
            with torch.autocast(enabled=False, device_type="cuda"):
                k = k.to(self.attn_dtype)
                v = v.to(self.attn_dtype)
        else:
            k, v = kv
            assert len(x) == len(k) and len(x) == len(v)

        # Attention
        out = _attn_func(q, k, v)

        # Output layer into (batch, seq_len, dim)
        out = out.flatten(2, 3)
        if self.use_gating:
            out = out * self.gate(x)
        out = self.fc(out)
        if self.use_sandwich:
            out = self.sandwich(out)

        return out, (k, v)

	import torch
	import torch.nn.functional as F
	import torch.utils.checkpoint

	# flash-attn is a hard dependency in the training / NVIDIA-container environment,
	# but it is unavailable (or ABI/arch-incompatible) on some inference targets such
	# as HF Spaces ZeroGPU (torch 2.8+, Blackwell sm_120). Soft-import it and fall
	# back to torch SDPA, which is numerically equivalent for our usage (non-causal,
	# no dropout, default scale = 1/sqrt(head_dim)).
	import os

	try:
	import flash_attn
	_HAS_FLASH = True
	except Exception: # ImportError or runtime/arch incompatibility on import
	flash_attn = None
	_HAS_FLASH = False

	# Allow forcing the SDPA path even when flash_attn is importable, so the
	# ZeroGPU code path can be validated on a flash-attn-equipped GPU node.
	if os.environ.get("DVSM_DISABLE_FLASH") == "1":
	_HAS_FLASH = False


	def _attn_func(q, k, v):
	"""Drop-in replacement for ``flash_attn.flash_attn_func``.

	flash_attn layout: q/k/v are ``[B, L, num_head, head_dim]``, non-causal,
	softmax_scale = 1/sqrt(head_dim). torch SDPA expects ``[B, num_head, L, head_dim]``
	and uses the same default scale + non-causal semantics, so we just transpose
	in and out. Output dtype matches the inputs (bf16/fp16), same as flash_attn.
	"""
	if _HAS_FLASH:
	return flash_attn.flash_attn_func(q, k, v)
	qh = q.transpose(1, 2) # [B, H, Lq, D]
	kh = k.transpose(1, 2) # [B, H, Lk, D]
	vh = v.transpose(1, 2) # [B, H, Lk, D]
	out = F.scaled_dot_product_attention(qh, kh, vh) # default scale, non-causal
	return out.transpose(1, 2).contiguous() # back to [B, Lq, H, D]


	def checkpoint_wrapper(module):
	# Wrap model forward function to automatically use activation checkpointing
	ori_forward = module.forward
	def new_forward(args, *kwargs):
	if module.training:
	output = torch.utils.checkpoint.checkpoint(
	ori_forward, args, *kwargs, use_reentrant=False)
	else:
	output = ori_forward(args, *kwargs)
	return output
	return new_forward


	class Contiguous(torch.nn.Module):
	def forward(self, x):
	return x.contiguous()


	class MLP(torch.nn.Module):
	def __init__(self, dim, expand_ratio=4, bias=False, activation=torch.nn.GELU):
	"""
	Args:
	dim: Input dimension
	expand_ratio: Multiplier for hidden dimension
	bias: Whether to use bias in linear layers
	activation: Activation function
	"""
	super().__init__()
	hidden_dim = round(dim * expand_ratio)
	self.mlp = torch.nn.Sequential(
	torch.nn.Linear(dim, hidden_dim, bias=bias),
	activation(),
	torch.nn.Linear(hidden_dim, dim, bias=bias))

	def forward(self, x):
	return self.mlp(x)


	class SwiGluMLP(torch.nn.Module):
	def __init__(self, dim, expand_ratio=4, bias=False):
	"""
	Args:
	dim: Input dimension
	expand_ratio: Multiplier for hidden dimension
	bias: Whether to use bias in linear layers
	activation: Activation function
	"""
	super().__init__()
	hidden_dim = round(dim * expand_ratio)
	self.mlp = torch.nn.Sequential(
	torch.nn.Linear(dim, hidden_dim, bias=bias),
	torch.nn.Linear(dim, hidden_dim, bias=bias),
	torch.nn.Linear(hidden_dim, dim, bias=bias))

	def forward(self, x):
	gate = torch.nn.functional.silu(self.mlp[0](x), inplace=True)
	return self.mlp[2](gate * self.mlp[1](x))


	class IntraViewAttention(torch.nn.Module):
	def __init__(self,
	dim,
	head_dim,
	bias=False,
	use_qk_norm=True,
	use_gating=False,
	use_sandwich=False,
	out_dim=None,
	):
	"""
	dim: Input dimension
	head_dim: Dimension of each attention head
	bias: Whether to use bias in QKV and output projection
	use_qk_norm: Whether to use Q-K normalization
	use_gating: Whether to apply gating before output ffn
	"""
	super().__init__()
	assert dim % head_dim == 0, f"dim={dim} should be divisible by head_dim={head_dim}"
	num_head = dim // head_dim

	self.dim = dim
	self.out_dim = dim if out_dim is None else out_dim
	self.head_dim = head_dim
	self.num_head = num_head
	self.use_qk_norm = use_qk_norm
	self.use_gating = use_gating
	self.use_sandwich = use_sandwich
	self.attn_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

	self.to_qkv = torch.nn.Linear(dim, 3 * dim, bias=bias)
	self.fc = torch.nn.Linear(dim, self.out_dim, bias=bias)

	# Optional Q-K normalization
	if self.use_qk_norm:
	self.q_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
	self.k_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
	else:
	self.q_norm = torch.nn.Identity()
	self.k_norm = torch.nn.Identity()

	# Optional gating
	if self.use_gating:
	self.gate = torch.nn.Sequential(
	torch.nn.Linear(dim, dim, bias=False),
	torch.nn.Sigmoid())
	# Optional normalization
	if self.use_sandwich:
	self.sandwich = torch.nn.LayerNorm(dim, bias=False)

	def forward(self, x):
	"""
	x: [B, L, D]
	"""
	# Predict qkv packed as (batch, seq_len, 3, num_head, head_dim)
	qkv = self.to_qkv(x).unflatten(2, (3, self.num_head, self.head_dim))

	# Unpack (q, k, v) each as (batch, seq_len, num_head, head_dim)
	q, k, v = qkv.moveaxis(2,0).contiguous()

	# Apply qk normalization if enabled
	q = self.q_norm(q)
	k = self.k_norm(k)

	# Attention
	with torch.autocast(enabled=False, device_type="cuda"):
	q = q.to(self.attn_dtype)
	k = k.to(self.attn_dtype)
	v = v.to(self.attn_dtype)

	out = _attn_func(q, k, v)

	# Output layer to yield (batch, seq_len, dim)
	out = out.flatten(2,3)
	if self.use_gating:
	out = out * self.gate(x)
	out = self.fc(out)
	if self.use_sandwich:
	out = self.sandwich(out)

	return out


	class CrossViewAttention(torch.nn.Module):
	def __init__(self,
	dim,
	head_dim,
	bias=False,
	use_qk_norm=True,
	use_gating=False,
	use_sandwich=False,
	out_dim=None,
	):
	"""
	Args:
	dim: Input dimension
	head_dim: Dimension of each attention head
	bias: Whether to use bias in QKV and output projection
	use_qk_norm: Whether to use Q-K normalization
	use_gating: Whether to apply gating before output ffn
	"""
	super().__init__()
	assert dim % head_dim == 0, f"dim={dim} should be divisible by head_dim={head_dim}"
	num_head = dim // head_dim

	self.dim = dim
	self.out_dim = dim if out_dim is None else out_dim
	self.head_dim = head_dim
	self.num_head = num_head
	self.use_qk_norm = use_qk_norm
	self.use_gating = use_gating
	self.use_sandwich = use_sandwich
	self.attn_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

	self.to_q = torch.nn.Linear(dim, dim, bias=bias)
	self.to_kv = torch.nn.Linear(dim, 2 * dim, bias=bias)
	self.fc = torch.nn.Linear(dim, self.out_dim, bias=bias)

	# Optional Q-K normalization
	if self.use_qk_norm:
	self.q_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
	self.k_norm = torch.nn.RMSNorm(head_dim, eps=1e-5)
	else:
	self.q_norm = torch.nn.Identity()
	self.k_norm = torch.nn.Identity()

	# Optional gating
	if self.use_gating:
	self.gate = torch.nn.Sequential(
	torch.nn.Linear(dim, dim, bias=False),
	torch.nn.Sigmoid())
	# Optional normalization
	if self.use_sandwich:
	self.sandwich = torch.nn.LayerNorm(dim, bias=False)

	def forward(self, x, kv=None):
	"""
	x: [B, L, D]
	kv: ([B, M, D], [B, M, D]) If provided, x query on kv instead of doing self attention
	"""
	q = self.to_q(x)

	# To shape of (batch, seq_len, num_head, head_dim)
	q = q.unflatten(2, (self.num_head, self.head_dim))
	q = self.q_norm(q)
	with torch.autocast(enabled=False, device_type="cuda"):
	q = q.to(self.attn_dtype)

	# Prepare (k, v)
	if kv is None:
	kv = self.to_kv(x)
	kv = kv.unflatten(2, (2, self.num_head, self.head_dim))
	k, v = kv.moveaxis(2, 0).contiguous()
	k = self.k_norm(k)
	with torch.autocast(enabled=False, device_type="cuda"):
	k = k.to(self.attn_dtype)
	v = v.to(self.attn_dtype)
	else:
	k, v = kv
	assert len(x) == len(k) and len(x) == len(v)

	# Attention
	out = _attn_func(q, k, v)

	# Output layer into (batch, seq_len, dim)
	out = out.flatten(2, 3)
	if self.use_gating:
	out = out * self.gate(x)
	out = self.fc(out)
	if self.use_sandwich:
	out = self.sandwich(out)

	return out, (k, v)