File size: 4,535 Bytes
df84732 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Prompt embedding cache: release text encoder (Qwen3.5) hidden states, one tensor per prompt.
torchrun --nproc-per-node N -m rapid_moss.textenc --model DIR --prompts train.jsonl eval.jsonl --out CACHE
Each rank writes cache/shard-<rank>-<k>.safetensors; rank 0 merges cache/index.json at the end.
Key "" is the empty negative prompt.
"""
from __future__ import annotations
import argparse
import json
import os
import time
from pathlib import Path
import torch
import torch.distributed as dist
from safetensors import safe_open
from safetensors.torch import save_file
def read_prompts(paths):
rows = []
for p in paths:
with open(p, encoding='utf-8') as f:
rows += [json.loads(line) for line in f if line.strip()]
return rows
class Encoder:
def __init__(self, model_dir, device):
from transformers import AutoModel, AutoTokenizer
self.tok = AutoTokenizer.from_pretrained(Path(model_dir) / 'tokenizer', local_files_only=True)
self.enc = AutoModel.from_pretrained(Path(model_dir) / 'text_encoder', dtype=torch.bfloat16,
local_files_only=True).to(device).eval()
self.device = device
@torch.no_grad()
def __call__(self, prompt, max_len=4096):
# Same as MOSSImagePipeline.encode_prompt for one prompt.
text = self.tok.apply_chat_template([{'role': 'user', 'content': prompt}], tokenize=False,
add_generation_prompt=True, enable_thinking=True)
inputs = self.tok(text, return_tensors='pt', truncation=True, max_length=max_len).to(self.device)
return self.enc(**inputs, use_cache=False, return_dict=True).last_hidden_state[0]
class Cache:
"""Read side: id -> [L, 4096] bf16 (CPU, memory-mapped shards)."""
def __init__(self, root):
self.root = Path(root)
self.index = json.loads((self.root / 'index.json').read_text())
self._files = {}
def __contains__(self, key):
return key in self.index
def get(self, key):
name = self.index[key]
if name not in self._files:
self._files[name] = safe_open(self.root / name, 'pt')
return self._files[name].get_tensor(key)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--model', required=True)
ap.add_argument('--prompts', nargs='+', required=True)
ap.add_argument('--out', required=True)
ap.add_argument('--shard-size', type=int, default=500)
args = ap.parse_args()
distributed = 'RANK' in os.environ
if distributed:
dist.init_process_group('nccl')
rank, world = (dist.get_rank(), dist.get_world_size()) if distributed else (0, 1)
device = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0)))
torch.cuda.set_device(device)
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
rows = read_prompts(args.prompts)
keys = [''] + [r['id'] for r in rows]
texts = dict([('', '')] + [(r['id'], r['prompt']) for r in rows])
if len(set(keys)) != len(keys):
raise SystemExit('duplicate prompt ids')
done = set()
for p in out.glob(f'shard-{rank}-*.safetensors'):
with safe_open(p, 'pt') as f:
done.update(f.keys())
mine = [k for i, k in enumerate(keys) if i % world == rank and k not in done]
enc = Encoder(args.model, device)
buf, shard, t0 = {}, len(list(out.glob(f'shard-{rank}-*.safetensors'))), time.time()
def flush():
nonlocal buf, shard
if buf:
save_file(buf, out / f'shard-{rank}-{shard:05d}.safetensors')
shard += 1
buf = {}
for n, key in enumerate(mine):
buf[key] = enc(texts[key]).to('cpu', torch.bfloat16).contiguous()
if len(buf) >= args.shard_size:
flush()
if rank == 0 and n % 200 == 0:
print(json.dumps({'encoded': n, 'of': len(mine), 's': round(time.time() - t0, 1)}), flush=True)
flush()
if distributed:
dist.barrier()
if rank == 0:
index = {}
for p in sorted(out.glob('shard-*.safetensors')):
with safe_open(p, 'pt') as f:
for k in f.keys():
index[k] = p.name
missing = [k for k in keys if k not in index]
(out / 'index.json').write_text(json.dumps(index))
print(json.dumps({'index': len(index), 'missing': len(missing)}), flush=True)
if distributed:
dist.destroy_process_group()
if __name__ == '__main__':
main()
|