| """Load MOSS-Image-SFT with OrbitQuant W4A8 components into the model's own pipeline (``pipeline_moss_image.py``).""" |
|
|
| from __future__ import annotations |
|
|
| import dataclasses |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
|
|
| GROUPS = ("qkv", "mlp", "attn_out", "mlp_out") |
|
|
|
|
| def _allocator(device: str) -> None: |
| |
| if not device.startswith("cuda") or {"PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF"} & set(os.environ): |
| return |
| if not torch.cuda.is_initialized(): |
| os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" |
| else: |
| setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None) |
| (setter or torch.cuda.memory._set_allocator_settings)("expandable_segments:True") |
|
|
|
|
| def _load_transformer(path: Path, transformer_class, device: str, fused: bool, attention: str, activation_bits: int): |
| import orbitquant |
| from orbitquant.adaln import RTNInt4Linear |
| from orbitquant.artifacts.manifest import OrbitQuantManifest |
| from orbitquant.config import OrbitQuantConfig |
| from orbitquant.layers import OrbitQuantLinear |
| from orbitquant.modeling import _parent_and_child, _set_child |
| from safetensors.torch import load_file |
|
|
| config = OrbitQuantConfig.from_dict(json.loads((path / "quantization_config.json").read_text())) |
| manifest = OrbitQuantManifest.from_dict(json.loads((path / "orbitquant_manifest.json").read_text())) |
| with torch.device("meta"): |
| model = transformer_class.from_config(transformer_class.load_config(str(path))) |
| for name in manifest.quantized_modules: |
| bits = manifest.module_bits.get(name) |
| layer_config = dataclasses.replace(config, weight_bits=bits) if bits and bits != config.weight_bits else config |
| parent, child = _parent_and_child(model, name) |
| _set_child(parent, child, OrbitQuantLinear.empty_from_linear( |
| model.get_submodule(name), config=layer_config, module_name=name)) |
| for name in manifest.adaln_modules: |
| parent, child = _parent_and_child(model, name) |
| _set_child(parent, child, RTNInt4Linear.empty_from_linear(model.get_submodule(name), config=config, module_name=name)) |
| model.load_state_dict(load_file(str(path / "model.safetensors"), device=device), strict=True, assign=True) |
| model = model.to(device).eval() |
| if fused: |
| from orbitquant.fused import fuse |
|
|
| widths = dict.fromkeys(GROUPS, activation_bits) if activation_bits != config.activation_bits else None |
| fuse(model, config=config, attention=attention, activation_bits=widths) |
| return model |
|
|
|
|
| def _park(model: nn.Module) -> None: |
| """Move a component to host memory right after its GPU build, so the next build does not stack on it.""" |
| from orbitquant.fused import release |
|
|
| release(model) |
| model.to("cpu") |
| torch.cuda.empty_cache() |
|
|
|
|
| class HostEmbedding(nn.Module): |
| """Token embedding kept in host memory: a prompt reads a few hundred of its 248k rows, so the 2 GiB table stays |
| off the GPU. The table is not a registered parameter, so moving the text encoder does not move it.""" |
|
|
| def __init__(self, embedding: nn.Embedding): |
| super().__init__() |
| self.__dict__["table"] = embedding.weight.detach().to("cpu") |
| self.num_embeddings, self.embedding_dim = self.table.shape |
|
|
| @property |
| def weight(self): |
| return self.table |
|
|
| def forward(self, ids: torch.Tensor) -> torch.Tensor: |
| return nn.functional.embedding(ids.cpu(), self.table).to(ids.device, non_blocking=True) |
|
|
|
|
| def _load_text_encoder(path: Path, device: str, fused: bool, activation_bits: int, host_embedding: bool): |
| import orbitquant |
| from transformers import AutoModel |
|
|
| model = AutoModel.from_pretrained(str(path), dtype=torch.bfloat16, device_map={"": device}).eval() |
| if fused: |
| from orbitquant.fused import fuse_linears |
|
|
| fuse_linears(model, activation_bits=activation_bits) |
| if host_embedding: |
| model.embed_tokens = HostEmbedding(model.embed_tokens) |
| return model |
|
|
|
|
| def load_pipeline(path: str | Path, device: str = "cuda", attention: str = "cudnn", offload: bool = False, |
| fused: bool | None = None, activation_bits: int = 8, host_embedding: bool = True): |
| """MOSS-Image pipeline with the W4A8 transformer and text encoder; call it as the original pipeline. |
| |
| ``attention`` is ``cudnn`` (default), ``flash``, ``sdpa`` or ``int8`` (faster, changes the image slightly). |
| ``offload`` keeps the components in host memory and moves each to the GPU only while it runs. ``fused=None`` |
| fuses the projections on CUDA devices with compute capability 8.0 or newer (needed for 8-bit activations).""" |
| path = Path(path) |
| _allocator(device) |
| if str(path) not in sys.path: |
| sys.path.insert(0, str(path)) |
| from diffusers import AutoencoderKLFlux2 |
| from pipeline_moss_image import MOSSImagePipeline, MOSSImageScheduler, MOSSImageTransformer |
| from transformers import AutoTokenizer |
|
|
| if fused is None: |
| fused = device.startswith("cuda") and torch.cuda.get_device_capability(device) >= (8, 0) |
| |
| |
| transformer = _load_transformer(path / "transformer", MOSSImageTransformer, device, fused, attention, activation_bits) |
| if offload: |
| _park(transformer) |
| text_encoder = _load_text_encoder(path / "text_encoder", device, fused, activation_bits, host_embedding) |
| if offload: |
| _park(text_encoder) |
| vae = AutoencoderKLFlux2.from_pretrained(str(path / "vae"), torch_dtype=torch.float32) |
| vae = vae.to("cpu" if offload else device).eval() |
| tokenizer = AutoTokenizer.from_pretrained(str(path / "tokenizer")) |
| pipeline = MOSSImagePipeline(transformer, MOSSImageScheduler(), text_encoder, tokenizer, vae) |
| if offload: |
| pipeline.enable_model_cpu_offload(device=device) |
| torch.cuda.empty_cache() |
| return pipeline |
|
|
