| |
| """Всё про MOSS-Image в одном месте: шаблон промпта, учитель (Qwen3.5-9B), ученик (0.8B), пайплайн. |
| |
| Задача та же, что в docs/qwen-image-2.1-adapter.md: заменить текстовый энкодер 15.6 ГБ |
| (Qwen3.5-VL-9B, `language_model`, hidden 4096) на замороженную 0.8B + адаптер. Но здесь |
| проще и одновременно строже: |
| |
| * токенизатор у учителя и ученика ОДИН (248 077 токенов, совпадают и строки, и id), |
| поэтому общая сегментация не нужна — id одни и те же; |
| * цель — `last_hidden_state`, то есть ПОСТ-норм выход слоя 32 (Qwen3_5TextModel сам |
| применяет `norm`). Это не то же самое, что в 2.1 (там был пред-норм выход), и именно |
| это видит DiT: text_embedder.proj 4096 -> 3072; |
| * учитель и ученик — одно семейство (Qwen3.5 hybrid linear attention), срезы ученика |
| берутся после его full-attention слоёв (3,7,11,15,19,23) -> hs[4,8,12,16,20,24]; |
| * пайплайн принимает `prompt_embeds`/`negative_prompt_embeds` официально, monkeypatch |
| не нужен (но метод `from_pretrained(text_encoder=None)` НЕ пропускает загрузку TE — |
| пайплайн собираем вручную, см. build_pipeline). |
| """ |
| import json |
| import os |
| import sys |
|
|
| import torch |
|
|
| MOSS_DIR = os.environ.get("MOSS_DIR", "/home/recoilme/moss") |
| STUDENT = os.environ.get("MOSS_STUDENT", "Qwen/Qwen3.5-0.8B") |
| |
| STUDENT_SUBDIR = "text_encoder" |
| TEACHER_SUBDIRS = ("text_encoder_qwen3.5-9b", "text_encoder") |
| |
| LAYERS = (4, 8, 12, 16, 20, 24) |
|
|
|
|
| |
| def load_tokenizer(moss_dir=None): |
| """Токенизатор учителя = токенизатор ученика, берём из папки модели.""" |
| from transformers import AutoTokenizer |
| return AutoTokenizer.from_pretrained(os.path.join(moss_dir or MOSS_DIR, "tokenizer")) |
|
|
|
|
| def render(tok, text): |
| """Ровно то, что делает `MOSSImagePipeline.encode_prompt` (сырая строка, не id).""" |
| return tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=True) |
|
|
|
|
| def encode(tok, texts, max_len=4096): |
| """list[str] -> list[list[int]]: шаблон + обрезка хвоста, как в пайплайне (truncation).""" |
| return [tok(render(tok, t), truncation=True, max_length=max_len)["input_ids"] for t in texts] |
|
|
|
|
| |
| def _has_config(path): |
| return os.path.isfile(os.path.join(path, "config.json")) |
|
|
|
|
| def teacher_dir(moss_dir=None): |
| """Папка оригинального текстового энкодера 9B (не удаляем и не переименовываем в text_encoder).""" |
| root = moss_dir or MOSS_DIR |
| for name in TEACHER_SUBDIRS: |
| path = os.path.join(root, name) |
| if _has_config(path): |
| return path |
| return os.path.join(root, TEACHER_SUBDIRS[0]) |
|
|
|
|
| def student_dir(moss_dir=None): |
| """Ученик: локальная копия в папке модели (`text_encoder/`), иначе HF-репозиторий.""" |
| local = os.path.join(moss_dir or MOSS_DIR, STUDENT_SUBDIR) |
| return local if _has_config(local) else STUDENT |
|
|
|
|
| def load_text_model(path, dtype=torch.bfloat16, load_8bit=False, device="cuda"): |
| """Qwen3.5 в любой раскладке — класс берём из `config.json`. |
| |
| Две раскладки одного и того же семейства: `Qwen3_5Model` (ключи `language_model.*` — так лежит |
| текстовый энкодер MOSS) и `Qwen3_5ForConditionalGeneration` (ключи `model.*` — родной чекпойнт |
| Qwen3.5-0.8B). Угадывать нельзя: не та раскладка грузится без ошибки, но со случайными весами |
| и только предупреждением в логе — поэтому смотрим architectures и падаем, если не узнали. |
| """ |
| import transformers |
| from transformers import AutoConfig |
| arch = (AutoConfig.from_pretrained(path, local_files_only=os.path.isdir(path)).architectures or [""])[0] |
| classes = {"Qwen3_5Model": transformers.Qwen3_5Model, |
| "Qwen3_5ForConditionalGeneration": transformers.Qwen3_5ForConditionalGeneration} |
| if arch not in classes: |
| raise SystemExit(f"{path}: не знаю, как грузить architectures={arch!r}") |
| local = {"local_files_only": os.path.isdir(path)} |
| if load_8bit: |
| from transformers import BitsAndBytesConfig |
| return classes[arch].from_pretrained( |
| path, quantization_config=BitsAndBytesConfig(load_in_8bit=True), |
| device_map={"": 0}, **local).eval() |
| return classes[arch].from_pretrained(path, dtype=dtype, **local).to(device).eval() |
|
|
|
|
| def text_stack(model): |
| """Текстовая стопка без vision-tower: у MOSS TE это `language_model`, у 0.8B — `model.language_model`.""" |
| for name in ("language_model", "model.language_model"): |
| obj = model |
| for part in name.split("."): |
| obj = getattr(obj, part, None) |
| if obj is None: |
| break |
| if obj is not None and hasattr(obj, "norm"): |
| return obj |
| raise SystemExit("не нашёл текстовую модель с .norm") |
|
|
|
|
| class Teacher: |
| """Qwen3.5-VL-9B (текстовая часть): цель = post-norm `last_hidden_state`, (B,L,4096).""" |
|
|
| def __init__(self, path=None, device="cuda", dtype=torch.bfloat16, load_8bit=False): |
| path = path or teacher_dir() |
| self.base = load_text_model(path, dtype=dtype, load_8bit=load_8bit, device=device) |
| self.text = text_stack(self.base) |
| self.quant = "8bit" if load_8bit else str(dtype).split(".")[-1] |
| cfg = self.text.config |
| print(f"учитель : {sum(p.numel() for p in self.base.parameters()) / 1e9:.2f}B, " |
| f"слоёв {cfg.num_hidden_layers}, hidden {cfg.hidden_size}, {self.quant}, " |
| f"VRAM {torch.cuda.memory_allocated() / 2**30:.1f} ГиБ", flush=True) |
|
|
| @torch.no_grad() |
| def last_hidden(self, ids, mask=None): |
| """(B,L) -> (B,L,4096) float32, пост-норм выход последнего слоя.""" |
| return self.text(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state.float() |
|
|
|
|
| def load_teacher(device="cuda", prefer_8bit=False, path=None): |
| """Учитель целиком, а если не влезает в карту — 8 бит. Возвращает (teacher, режим). |
| |
| После неудачной попытки надо отпустить traceback: он держит кадр с недогруженной моделью, |
| и без gc/empty_cache вторая попытка падает на уже занятой памяти (проверено). |
| """ |
| order = [True, False] if prefer_8bit else [False, True] |
| errors = [] |
| for load_8bit in order: |
| name = "8-бит" if load_8bit else "bf16" |
| try: |
| return Teacher(path=path, device=device, load_8bit=load_8bit), \ |
| ("8bit" if load_8bit else "bf16") |
| except (torch.OutOfMemoryError, RuntimeError) as exc: |
| if "out of memory" not in str(exc).lower(): |
| raise |
| free_gib = torch.cuda.mem_get_info()[0] / 2**30 |
| print(f" {name} не влез ({free_gib:.1f} ГиБ свободно) — пробую дальше", flush=True) |
| errors.append(f"{name}: {str(exc).splitlines()[0][:100]}") |
| del exc |
| empty_cache() |
| raise SystemExit("учитель не грузится: " + " | ".join(errors)) |
|
|
|
|
| class Student: |
| """Qwen3.5-0.8B (заморожен): срезы hidden_states -> (B, L, K*hidden).""" |
|
|
| def __init__(self, path=None, layers=LAYERS, device="cuda", dtype=torch.bfloat16): |
| path = path or student_dir() |
| self.path = path |
| self.base = load_text_model(path, dtype=dtype, device=device) |
| self.base.requires_grad_(False) |
| self.text = text_stack(self.base) |
| if self.text.config.hidden_size >= 4096: |
| raise SystemExit(f"{path}: hidden {self.text.config.hidden_size} — это учитель, а не " |
| f"ученик (у Qwen3.5-0.8B hidden 1024)") |
| self.layers = list(layers) |
| self.device = device |
| cfg = self.text.config |
| print(f"ученик : {sum(p.numel() for p in self.base.parameters()) / 1e9:.2f}B, " |
| f"слоёв {cfg.num_hidden_layers}, hidden {cfg.hidden_size}, срезы {self.layers}", |
| flush=True) |
|
|
| @property |
| def in_dim(self): |
| return len(self.layers) * self.text.config.hidden_size |
|
|
| @torch.no_grad() |
| def hidden(self, ids, mask=None): |
| out = self.text(input_ids=ids, attention_mask=mask, output_hidden_states=True, use_cache=False) |
| hs = out.hidden_states |
| assert hs is not None and len(hs) > max(self.layers), f"скрытых состояний {len(hs)}" |
| x = torch.stack([hs[i] for i in self.layers], dim=1) |
| return x.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[2], -1).float() |
|
|
|
|
| |
| def build_pipeline(moss_dir=None, dit_dtype=torch.bfloat16, vae_dtype=torch.float32): |
| """MOSSImagePipeline без текстового энкодера. |
| |
| `DiffusionPipeline.from_pretrained(..., text_encoder=None)` здесь НЕ работает: text_encoder |
| объявлен в __init__ с дефолтом None, поэтому diffusers не считает его «переданным» и лезет |
| грузить 15.6 ГБ. Собираем компоненты руками — их три, и все локальные. |
| """ |
| moss_dir = moss_dir or MOSS_DIR |
| if moss_dir not in sys.path: |
| sys.path.insert(0, moss_dir) |
| from diffusers import AutoencoderKLFlux2 |
| from pipeline_moss_image import MOSSImagePipeline, MOSSImageScheduler, MOSSImageTransformer |
|
|
| transformer = MOSSImageTransformer.from_pretrained( |
| os.path.join(moss_dir, "transformer"), torch_dtype=dit_dtype, local_files_only=True) |
| vae = AutoencoderKLFlux2.from_pretrained( |
| os.path.join(moss_dir, "vae"), torch_dtype=vae_dtype, local_files_only=True) |
| cfg = json.load(open(os.path.join(moss_dir, "scheduler", "scheduler_config.json"))) |
| keys = ("num_train_timesteps", "shift", "training_shift") |
| scheduler = MOSSImageScheduler(**{k: cfg[k] for k in keys if k in cfg}) |
| return MOSSImagePipeline(transformer.eval(), scheduler, text_encoder=None, tokenizer=None, |
| vae=vae.eval()) |
|
|
|
|
| @torch.no_grad() |
| def condition(adapter, student, tok, texts, max_len=4096): |
| """list[str] -> list[(embeds [L,4096] bf16, mask [L] bool)]: по одному промпту без паддинга. |
| |
| Пайплайн кодирует каждый промпт отдельно (`unpadded individual forwards`), поэтому и здесь |
| так же: длину маскирует сам DiT, а не токенизатор. |
| """ |
| device = next(adapter.parameters()).device |
| out = [] |
| for ids in encode(tok, texts, max_len): |
| si = torch.tensor([ids], dtype=torch.long, device=device) |
| mask = torch.ones_like(si, dtype=torch.bool) |
| y = adapter(student.hidden(si, mask), mask) |
| out.append((y[0].to(torch.bfloat16).cpu(), mask[0].cpu())) |
| return out |
|
|
|
|
| @torch.no_grad() |
| def generate(pipe, pos, neg, size=1024, steps=30, seed=1234, guidance=3.0, device="cuda"): |
| """pos/neg: (embeds [L,4096], mask [L]) -> PIL. CFG включён всегда (uncond_skip_layers).""" |
| emb, mask = pos |
| nemb, nmask = neg |
| gen = torch.Generator(device).manual_seed(seed) |
| return pipe(prompt_embeds=emb[None].to(device), prompt_attention_mask=mask[None].to(device), |
| negative_prompt_embeds=nemb[None].to(device), |
| negative_prompt_attention_mask=nmask[None].to(device), |
| height=size, width=size, num_inference_steps=steps, guidance_scale=guidance, |
| generator=gen, output_type="pil").images[0] |
|
|
|
|
| def empty_cache(): |
| """Вернуть освобождённую память карты. Сами объекты надо удалить у себя (`del`): `del` внутри |
| функции снимает только локальную ссылку, объект остаётся жив, пока на него смотрит вызывающий.""" |
| import gc |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
