| |
| """Strict-limit RAM sweep for EPQ on LFM2.5-24B-A2B Q4_K_M. |
| |
| Unlike the original memory_levels_ling.py (which only sets the pager's own |
| budget), every run here is launched under a *hard* OS limit: |
| |
| prlimit --data=<level> --nofile=65536 |
| |
| On Linux >= 4.7 RLIMIT_DATA covers brk() AND mmap'd anonymous memory, so it is |
| a real ceiling on the arena + heap + every later anonymous allocation. If the |
| process exceeds it, malloc fails and the run aborts -- we record that as a hard |
| failure instead of a silent overshoot. |
| |
| Two caveats this harness records explicitly: |
| |
| * RLIMIT_DATA polices *anonymous* memory only. The ~1.3-2.1 GiB of |
| non-expert weights are mmap'd file-backed pages and are therefore NOT |
| policed here, so peak RSS can exceed the level. Both numbers are reported. |
| * `-march=native` on this host defines __AMX_INT8__, which makes llama.cpp |
| allocate a ~1.88 GiB AMX repack buffer before the arena-sized context. That |
| alone pushes the floor past 4 GiB, so the sweep runs against the |
| GGML_NATIVE=OFF (no-AMX) build unless --amx is passed. |
| |
| Swap is proven absent at run time by sampling /proc/<pid>/status VmSwap; a |
| nonzero value fails the run. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import platform |
| import re |
| import statistics |
| import subprocess |
| import sys |
| import threading |
| import time |
| from pathlib import Path |
|
|
| MIB = 1024 * 1024 |
| GIB = 1024 ** 3 |
|
|
| PROMPT = ("Write a detailed paragraph about how a mixture-of-experts model routes " |
| "tokens to expert weights, and why paging those weights from SSD matters.") |
|
|
|
|
| def args_config(args): |
| """JSON-safe view of the argparse namespace (drops the resolve helper).""" |
| return {k: v for k, v in vars(args).items() if not callable(v)} |
|
|
|
|
| def host_facts() -> dict: |
| cpu = "" |
| for line in Path('/proc/cpuinfo').read_text().splitlines(): |
| if line.startswith('model name'): |
| cpu = line.split(':', 1)[1].strip() |
| break |
| mem_total = 0 |
| for line in Path('/proc/meminfo').read_text().splitlines(): |
| if line.startswith('MemTotal'): |
| mem_total = int(line.split()[1]) * 1024 |
| break |
| cg = {} |
| for k, f in (('memory.max', 'memory.max'), |
| ('memory.swap.max', 'memory.swap.max'), |
| ('memory.high', 'memory.high')): |
| try: |
| cg[k] = Path('/sys/fs/cgroup/' + f).read_text().strip() |
| except Exception: |
| cg[k] = None |
| return {'cpu_model': cpu, 'cpu_threads': os.cpu_count(), |
| 'mem_total_bytes': mem_total, 'cgroup': cg} |
|
|
|
|
| class ProcSampler(threading.Thread): |
| """Sample RSS / AnonHWM / VmSwap until told to stop. |
| |
| AnonHWM is smaps_rollup's own anonymous high-water mark and is the quantity |
| RLIMIT_DATA actually polices, so it is what we compare against the level. |
| """ |
|
|
| def __init__(self, pid: int, interval: float = 0.03): |
| super().__init__(daemon=True) |
| self.pid = pid |
| self.interval = interval |
| self.stop_flag = threading.Event() |
| self.peak_rss = 0 |
| self.peak_anon = 0 |
| self.peak_swap = 0 |
| self.n = 0 |
|
|
| def _read(self) -> tuple[int, int, int]: |
| rss = anon = swp = 0 |
| try: |
| for line in Path(f'/proc/{self.pid}/status').read_text().splitlines(): |
| if line.startswith('VmRSS:'): |
| rss = int(line.split()[1]) * 1024 |
| elif line.startswith('VmSwap:'): |
| swp = int(line.split()[1]) * 1024 |
| except OSError: |
| pass |
| try: |
| for line in Path(f'/proc/{self.pid}/smaps_rollup').read_text().splitlines(): |
| if line.startswith('Rss:'): |
| rss = max(rss, int(line.split()[1]) * 1024) |
| elif line.startswith('Anonymous:'): |
| anon = int(line.split()[1]) * 1024 |
| elif line.startswith('AnonHWM:'): |
| anon = max(anon, int(line.split()[1]) * 1024) |
| except OSError: |
| pass |
| return rss, anon, swp |
|
|
| def run(self): |
| while not self.stop_flag.is_set(): |
| r, a, s = self._read() |
| self.n += 1 |
| self.peak_rss = max(self.peak_rss, r) |
| self.peak_anon = max(self.peak_anon, a) |
| self.peak_swap = max(self.peak_swap, s) |
| time.sleep(self.interval) |
| r, a, s = self._read() |
| self.peak_rss = max(self.peak_rss, r) |
| self.peak_anon = max(self.peak_anon, a) |
| self.peak_swap = max(self.peak_swap, s) |
|
|
|
|
| def evict(shards: list[str], helper: str) -> None: |
| subprocess.run(['sync'], check=False) |
| if helper: |
| subprocess.run([helper, '--evict'] + shards, check=False) |
|
|
|
|
| def run_once(args, level_mb: int, repeat: int) -> dict: |
| stats_file = Path(args.output_dir) / f'level_{level_mb}_r{repeat}.jsonl' |
| log_file = Path(args.output_dir) / f'level_{level_mb}_r{repeat}.log' |
| stats_file.unlink(missing_ok=True) |
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| inner = [ |
| str(Path(args.bin).resolve()), |
| '-m', str(Path(args.model).resolve()), |
| '-p', PROMPT, '-n', str(args.n_predict), '-c', str(args.ctx), |
| '-t', str(args.threads), '-b', str(args.ubatch), '--temp', '0', |
| '--no-interactive', '--no-chat', '--stats-file', str(stats_file), |
| '--ram-budget-mb', str(level_mb), |
| ] |
| |
| cmd = ['prlimit', f'--data={level_mb * MIB}', '--nofile=65536'] + inner |
|
|
| env = dict(os.environ) |
| env.update({ |
| 'SDQ_RAM_BUDGET_MB': str(level_mb), |
| 'SDQ_RESERVE_MB': str(args.reserve_mb_for(level_mb)), |
| 'SDQ_PREFETCH': '0', |
| 'SDQ_FADV_DONTNEED': '1', |
| 'SDQ_BENCH_NO_EOG': '1', |
| }) |
|
|
| evict(args.shards, args.evict_helper) |
|
|
| t0 = time.perf_counter() |
| proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, |
| text=True, encoding='utf-8', errors='replace', env=env) |
| sampler = ProcSampler(proc.pid) |
| sampler.start() |
| try: |
| out, err = proc.communicate(timeout=args.timeout) |
| except subprocess.TimeoutExpired: |
| proc.kill() |
| out, err = proc.communicate() |
| sampler.stop_flag.set() |
| return {'level_mb': level_mb, 'repeat': repeat + 1, |
| 'error': f'timeout after {args.timeout}s'} |
| sampler.stop_flag.set() |
| sampler.join(timeout=2) |
| wall = time.perf_counter() - t0 |
| both = out + '\n' + err |
| log_file.write_text(both, encoding='utf-8', errors='replace') |
|
|
| res = { |
| 'level_mb': level_mb, |
| 'level_gib': level_mb / 1024, |
| 'repeat': repeat + 1, |
| 'wall_s': round(wall, 2), |
| 'exit_code': proc.returncode, |
| 'peak_rss_mb': round(sampler.peak_rss / MIB, 1), |
| 'peak_anon_mb': round(sampler.peak_anon / MIB, 1), |
| 'peak_swap_mb': round(sampler.peak_swap / MIB, 1), |
| 'swap_free': sampler.peak_swap == 0, |
| 'anon_within_limit': sampler.peak_anon <= level_mb * MIB, |
| 'limit_enforced_by': 'prlimit --data (RLIMIT_DATA)', |
| } |
|
|
| if proc.returncode != 0: |
| sig = -proc.returncode if proc.returncode < 0 else None |
| res['signal'] = sig |
| res['status'] = 'OOM_OR_ABORT' |
| res['failure'] = ' | '.join( |
| [l.strip() for l in err.splitlines() |
| if re.search(r'bad_alloc|cannot allocate|MemoryError|std::terminate|' |
| r'abort|Killed|Expert arena|too small|error', l, re.I)][-5:]) |
| res['wall_s'] = round(wall, 2) |
| return res |
|
|
| res['status'] = 'OK' |
|
|
| def grab(pat, text=both): |
| m = re.search(pat, text, re.M) |
| return m |
|
|
| d = grab(r'decode:\s*(\d+)\s+tok\s*/\s*([\d.]+)s\s*=\s*([\d.]+)\s+tok/s') |
| p = grab(r'prefill:\s*(\d+)\s+tok\s*/\s*([\d.]+)s\s*=\s*([\d.]+)\s+tok/s') |
| mem = grab(r'總 RSS\s+(\d+)\s+MB\s*/\s+峰值\s+(\d+)\s+MB\s*/\s+swap\s+(\d+)\s+MB') |
| split = grab(r'匿名 RSS\s+(\d+)\s+MB\s*/\s+檔案對映 RSS\s+(\d+)\s+MB\s*/\s+總 RSS\s+(\d+)\s+MB\s*/\s+峰值\s+(\d+)\s+MB') |
| arena = grab(r'slots=\d+\s+arena=([\d.]+)\s+GiB') |
| slots = grab(r'slots=(\d+)') |
| ssd_tot = grab(r'全程 SSD 實際讀取:\s*([\d.]+)\s*GB') |
| mib = grab(r'SSD 讀\s*([\d.]+)\s*MB\s*\(([\d.]+)\s*MB/tok\)') |
|
|
| if d: |
| res['decode_tok_s'] = float(d.group(3)) |
| res['decode_tok'] = int(d.group(1)) |
| if p: |
| res['prefill_tok_s'] = float(p.group(3)) |
| res['prefill_tok'] = int(p.group(1)) |
| if mem: |
| res['final_rss_mb'] = int(mem.group(1)) |
| res['selfpeak_rss_mb'] = int(mem.group(2)) |
| res['self_swap_mb'] = int(mem.group(3)) |
| if split: |
| res['rss_anon_mb'] = int(split.group(1)) |
| res['rss_file_mb'] = int(split.group(2)) |
| res['rss_total_mb'] = int(split.group(3)) |
| res['selfpeak_rss_mb'] = int(split.group(4)) |
| if arena: |
| res['arena_gib'] = float(arena.group(1)) |
| if slots: |
| res['n_slots'] = int(slots.group(1)) |
| if ssd_tot: |
| res['total_ssd_gb'] = float(ssd_tot.group(1)) |
| if mib: |
| res['decode_ssd_mb'] = float(mib.group(1)) |
| res['ssd_mb_per_token'] = float(mib.group(2)) |
|
|
| rec = {} |
| if stats_file.exists(): |
| for line in stats_file.read_text().splitlines(): |
| r = json.loads(line) |
| if r.get('phase') == 'decode': |
| rec = r |
| stats_file.unlink(missing_ok=True) |
|
|
| pl = rec.get('per_layer') or [] |
| rates = [x['hit_rate'] for x in pl if x.get('req', 0) > 0] |
| hr = rec.get('hit_rate') |
| res.update({ |
| 'expert_requests': rec.get('expert_requests'), |
| 'expert_hits': rec.get('expert_hits'), |
| 'hot_rate_pct': hr, |
| 'cold_rate_pct': (100.0 - hr) if hr is not None else None, |
| 'evictions': rec.get('evictions'), |
| 'per_layer_hit_min': min(rates) if rates else None, |
| 'per_layer_hit_max': max(rates) if rates else None, |
| }) |
|
|
| asst = '' |
| started = False |
| for line in out.splitlines(): |
| if line.startswith('assistant:'): |
| started = True |
| rest = line[len('assistant:'):].strip() |
| if rest: |
| asst = rest |
| continue |
| if started and line.strip() and not line.startswith('['): |
| asst += ' ' + line.strip() |
| res['assistant'] = asst.strip()[:200] |
| return res |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--model', required=True) |
| ap.add_argument('--shards', nargs='+', required=True) |
| ap.add_argument('--bin', required=True) |
| ap.add_argument('--levels-mb', type=int, nargs='+', default=[4096, 8192, 16384, 32768]) |
| ap.add_argument('--repeat', type=int, default=2) |
| ap.add_argument('--ctx', type=int, default=2048) |
| ap.add_argument('--ubatch', type=int, default=128) |
| ap.add_argument('--threads', type=int, default=16) |
| ap.add_argument('--n-predict', type=int, default=48) |
| ap.add_argument('--reserve-mb', type=int, default=2400) |
| ap.add_argument('--timeout', type=int, default=1800) |
| ap.add_argument('--evict-helper', default='/tmp/pagecache_resident') |
| ap.add_argument( |
| '--reserve-map', |
| default='4096:2400,8192:2400,16384:2400,32768:2400', |
| help='per-level SDQ_RESERVE_MB, "level_mb:reserve_mb,...". Calibrated ' |
| 'by measuring anon_hwm - arena, i.e. every anonymous allocation ' |
| 'that happens after the arena is sized.') |
| ap.add_argument('--output-dir', default='/root/app/ling/results_strict') |
| ap.add_argument('--out-json', default='/root/app/ling/results_strict/strict_memory_levels.json') |
| args = ap.parse_args() |
|
|
| args._reserve = {} |
| for item in args.reserve_map.split(','): |
| k, v = item.split(':') |
| args._reserve[int(k)] = int(v) |
| args.reserve_mb_for = lambda lvl: args._reserve.get(lvl, 2400) |
|
|
| Path(args.output_dir).mkdir(parents=True, exist_ok=True) |
|
|
| print(f'host: {host_facts()}', file=sys.stderr) |
| runs = [] |
| for lvl in args.levels_mb: |
| for r in range(args.repeat): |
| print(f'--- level {lvl} MB repeat {r + 1}/{args.repeat} ---', |
| file=sys.stderr, flush=True) |
| try: |
| res = run_once(args, lvl, r) |
| except Exception as e: |
| res = {'level_mb': lvl, 'repeat': r + 1, 'error': repr(e)} |
| runs.append(res) |
| print(json.dumps(res, ensure_ascii=False)[:400], file=sys.stderr, flush=True) |
| Path(args.out_json).write_text(json.dumps( |
| {'host': host_facts(), 'config': args_config(args), 'runs': runs}, |
| indent=2, ensure_ascii=False)) |
|
|
| print('DONE', file=sys.stderr) |
| Path(args.out_json).write_text(json.dumps( |
| {'host': host_facts(), 'config': args_config(args), 'runs': runs}, |
| indent=2, ensure_ascii=False)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
