#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""s27_patch_modeling.py — 修 modeling_samai_27b.py 跨设备聚合 bug
(ponder KL / MoE aux / moe_active_mean 的 torch.stack 混合 cuda:0+cuda:1)
幂等: 已打补丁则跳过"""
import glob, shutil

F = glob.glob("/tmp/k27/hf_home/models--tchbcb--samai-27b/snapshots/*/modeling_samai_27b.py")[0]
src = open(F).read()
orig = src

def rep(old, new):
    global src
    if new in src:
        print("skip (already):", old.strip()[:50])
        return
    assert old in src, f"NOT FOUND: {old[:70]!r}"
    src = src.replace(old, new, 1)
    print("patched:", old.strip()[:60])

# 模块级 helper (追加到文件尾)
if "_stack_dev" not in src:
    src += '''

def _stack_dev(ts):
    """跨设备安全聚合: 统一到首个张量的设备再 stack.mean()"""
    if not ts:
        return None
    ts = [t.to(ts[0].device) for t in ts]
    return torch.stack(ts).mean()
'''
    print("appended: _stack_dev helper")

# R1: _ponder_kl 聚合
rep("        return torch.stack(kls).mean() if kls else None",
    "        kls = [k.to(kls[0].device) for k in kls]\n"
    "        return torch.stack(kls).mean() if kls else None")

# R2: _moe_aux 聚合
rep("        return torch.stack(auxs).mean()",
    "        auxs = [a.to(auxs[0].device) for a in auxs]\n"
    "        return torch.stack(auxs).mean()")

# R3: moe_active_mean (输出指标, 但求值于 return 内, 必须修)
rep("""            moe_active_mean=(torch.stack([pl.mlp.last_active for pl in self.ponder_layers
                                          if pl.mlp.last_active is not None]).mean()
                             if any(pl.mlp.last_active is not None
                                    for pl in self.ponder_layers) else None),""",
    "            moe_active_mean=_stack_dev([pl.mlp.last_active for pl in self.ponder_layers\n"
    "                                        if pl.mlp.last_active is not None]),")

# R4/R5/R6: 损失相加设备对齐
rep("loss = loss + float(cfg.ponder_loss_beta) * ponder_kl",
    "loss = loss + float(cfg.ponder_loss_beta) * ponder_kl.to(loss.device)")
rep("loss = loss + float(cfg.moe_aux_weight) * moe_aux",
    "loss = loss + float(cfg.moe_aux_weight) * moe_aux.to(loss.device)")
rep("loss = loss + float(cfg.mtp_loss_weight) * mtp_loss",
    "loss = loss + float(cfg.mtp_loss_weight) * mtp_loss.to(loss.device)")

if src != orig:
    shutil.copy(F, F + ".bak_devpatch")
    open(F, "w").write(src)
    print("WROTE", F, "(backup .bak_devpatch)")
else:
    print("NO CHANGE")
import py_compile
py_compile.compile(F, doraise=True)
print("SYNTAX OK")