From e061384aa41df3b6b0ab68a1df415a8133cd17d2 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sun, 5 Jul 2026 14:18:05 -0400 Subject: [PATCH 1/6] fix(mp): FLOP-weighted calibration budget, per-row qk MP, drop alpha/beta - calibrate_mp_thresholds.py: budget is now sum(rows * macs * stoc_len) instead of sum(rows * stoc_len), so an "avgX" target means X in real compute, not row count. macs_per_row (out*in for linears, N*head_dim for qk/av) is threaded through add() -> the global lambda solve + iso-budget check. qk recording is switched to per-row. Validated: realized FLOP-weighted avg = 47.97 for target 48. - sc_attention.py: _sc_qk MP is now per-query-row (||Q_row||_inf), matching the per_row quant scale; removed the per-head classifier and static per-head path. - quant_sc_main.py: removed all alpha/beta remnants -- --mp_alpha/--mp_beta and the 14 per-op variants, the crashing AdaptiveMPConfig(alpha=,beta=, operator_params=) construction, and the run-name tag. --adaptive_mp now requires --adaptive_mp_table. - sc_attention.py + sc_mlp.py: repaired the adaptive_classify_rows signature drift vs kernel 9358239 (6 call sites passed the old arg order -> TypeError, so the adaptive path could not run at all). Invalidates existing calibration threshold tables (new cost model -> new allocations); re-calibration required. Co-Authored-By: Claude Opus 4.8 (1M context) --- qdit/sc_integration/sc_attention.py | 134 +++++++++++++--------------- qdit/sc_integration/sc_mlp.py | 8 +- scripts/calibrate_mp_thresholds.py | 97 +++++++++++++++----- scripts/quant_sc_main.py | 126 +++----------------------- 4 files changed, 156 insertions(+), 209 deletions(-) diff --git a/qdit/sc_integration/sc_attention.py b/qdit/sc_integration/sc_attention.py index e12c914..65f31b5 100644 --- a/qdit/sc_integration/sc_attention.py +++ b/qdit/sc_integration/sc_attention.py @@ -321,12 +321,12 @@ def _sc_linear_dynamic_mp(self, x, weight, bias, operator, chunk_d=0, grouped=Fa if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config @@ -426,12 +426,12 @@ def _sc_linear_combined_mp(self, x, weight, bias, operator, dispatch, if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config @@ -645,95 +645,85 @@ def _sc_linear(self, x, weight, bias, operator=None, chunk_d=0, # ================================================================= def _sc_qk(self, q, k): - """SC q@k^T matmul — supports per-head mixed precision.""" + """SC q@k^T matmul — per-query-row mixed precision. + + MP granularity is per query-row (per (batch, head)), consistent with the + per_row quantization scale. Rows are classified by ``||Q_row||_inf``. The + old per-head stoc_len path was removed so MP and quant granularity match. + """ q_scaled = q * self.scale B, H, N, D = q_scaled.shape - head_stoc_lens = self.sc_controller.get_group_stoc_lens( - self.block_idx, "qk") + has_dynamic_mp = (self.sc_controller.adaptive_mp_config is not None + or self.sc_controller.mp_config is not None) + + if not has_dynamic_mp: + # Uniform precision — existing fast path (batched kernel). + stoc_len = self.sc_controller.get_stoc_len(self.block_idx, "qk") + sc_prec = self.sc_controller.resolve_sc_prec(stoc_len) + return self._sc_qk_uniform(q_scaled, k, sc_prec, stoc_len) - # Dynamic MP: compute per-head stoc_lens from Q magnitude - if head_stoc_lens is None and (self.sc_controller.adaptive_mp_config is not None - or self.sc_controller.mp_config is not None): - q_metric = q_scaled.float().abs().amax(dim=(0, 2, 3)) # [H] - MetricProfiler.record(q_metric, self.sc_controller.current_timestep, - self.block_idx, "qk") + # Dynamic per-query-row MP (adaptive or fixed): one stoc_len per query + # row within each (batch, head). Mirrors _sc_av's per-row path and keeps + # MP granularity aligned with the per_row quant scale. + BH = B * H + q_flat = q_scaled.reshape(BH, N, D).float() + k_flat = k.reshape(BH, N, D).float() + output = torch.zeros(BH, N, N, device=q.device, dtype=torch.float32) + + matmul_fn = self._get_matmul_fn() + baseline_stoc_len = self.sc_controller.stoc_len + compute_baseline = 0 + compute_actual = 0.0 + + for i in range(BH): + row_metric = q_flat[i].abs().amax(dim=-1) # [N] = ||Q_row||_inf + if i == 0: # profile / log once per (t, block) + MetricProfiler.record(row_metric, self.sc_controller.current_timestep, + self.block_idx, "qk") if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( - q_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, + row_metric, self.sc_controller.adaptive_mp_config, operator="qk", block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config assignment = classify_rows_by_metric( - q_metric, mp_config.stoc_len_levels, mp_config.level_fractions) - - MPDistributionLogger.log( - self.sc_controller.current_timestep, self.block_idx, - "qk", assignment, H) - - head_stoc_lens = [0] * H - for sl, heads in assignment.level_row_indices.items(): - for h_idx in heads: - head_stoc_lens[h_idx.item()] = sl + row_metric, mp_config.stoc_len_levels, mp_config.level_fractions) - if head_stoc_lens is None: - # Uniform precision — existing fast path - stoc_len = self.sc_controller.get_stoc_len(self.block_idx, "qk") - sc_prec = self.sc_controller.resolve_sc_prec(stoc_len) - return self._sc_qk_uniform(q_scaled, k, sc_prec, stoc_len) - - # Mixed: group heads by stoc_len, compute within each group - output = torch.zeros(B, H, N, N, device=q.device, dtype=torch.float32) - stoc_len_to_heads: dict[int, list[int]] = defaultdict(list) - for h, sl in enumerate(head_stoc_lens): - stoc_len_to_heads[sl].append(h) - - baseline_stoc_len = self.sc_controller.stoc_len - compute_baseline = B * H * N * N * D * baseline_stoc_len - compute_actual = 0.0 + if i == 0: + MPDistributionLogger.log( + self.sc_controller.current_timestep, self.block_idx, + "qk", assignment, N) - for sl, heads in stoc_len_to_heads.items(): - # Pruned heads (stoc_len=0): output already zeroed - if sl == 0: - continue - compute_actual += B * len(heads) * N * N * D * sl + k_i = k_flat[i] # [N, D] + for sl, rows in assignment.level_row_indices.items(): + if len(rows) == 0 or sl == 0: + continue # pruned rows: output already zeroed + n_rows = len(rows) + compute_baseline += n_rows * N * D * baseline_stoc_len + compute_actual += n_rows * N * D * sl - sp = self.sc_controller.resolve_sc_prec(sl) - config = self._get_sc_config(D, sp) - - for h in heads: - # Extract single head: [B, 1, N, D] -> [B, N, D] - q_h = q_scaled[:, h].float() - k_h = k[:, h].float() - - matmul_fn = self._get_matmul_fn() - if self.sc_mode == "bipolar": - output[:, h] = matmul_fn( - q_h, k_h, granularity="per_head", mode="bipolar", - sc_prec=sp, config=config, stoc_len=sl, - rng_levels=self._rng_levels(sl), - ) - else: - for b_idx in range(B): - output[b_idx, h] = matmul_fn( - q_h[b_idx], k_h[b_idx], - granularity="per_tensor", mode=self.sc_mode, - sc_prec=sp, config=config, stoc_len=sl, - rng_levels=self._rng_levels(sl), - ) + sp = self.sc_controller.resolve_sc_prec(sl) + config = self._get_sc_config(D, sp) + q_sub = q_flat[i][rows] # [n_rows, D] + output[i, rows] = matmul_fn( + q_sub, k_i, + granularity="per_row", group_a=1, group_b=1, + mode=self.sc_mode, sc_prec=sp, config=config, + stoc_len=sl, rng_levels=self._rng_levels(sl)) MPDistributionLogger.log_compute( self.sc_controller.current_timestep, self.block_idx, "qk", compute_baseline, compute_actual) - return output + return output.reshape(B, H, N, N) def _sc_qk_uniform(self, q_scaled, k, sc_prec, stoc_len): """Uniform precision QK — fully batched kernel or per-head loop.""" @@ -852,12 +842,12 @@ def _sc_av(self, attn, v): if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_max, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator="av", block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config diff --git a/qdit/sc_integration/sc_mlp.py b/qdit/sc_integration/sc_mlp.py index 031198b..215936a 100644 --- a/qdit/sc_integration/sc_mlp.py +++ b/qdit/sc_integration/sc_mlp.py @@ -209,12 +209,12 @@ def _sc_linear_dynamic_mp(self, x, weight, bias, operator, chunk_d=0): if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config @@ -295,12 +295,12 @@ def _sc_linear_combined_mp(self, x, weight, bias, operator, dispatch, if self.sc_controller.adaptive_mp_config is not None: assignment = adaptive_classify_rows( row_metric, - self.sc_controller.current_timestep, - self.sc_controller.total_timesteps, self.sc_controller.adaptive_mp_config, operator=operator, block_idx=self.block_idx, total_blocks=self.sc_controller.total_blocks, + timestep=self.sc_controller.current_timestep, + total_timesteps=self.sc_controller.total_timesteps, ) else: mp_config = self.sc_controller.mp_config diff --git a/scripts/calibrate_mp_thresholds.py b/scripts/calibrate_mp_thresholds.py index c301dfe..4b70425 100644 --- a/scripts/calibrate_mp_thresholds.py +++ b/scripts/calibrate_mp_thresholds.py @@ -145,6 +145,7 @@ def _cosine_dist_heads(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor _METRIC_ROWS = _relative_l2_rows _METRIC_HEADS = _relative_l2_heads _USE_FP_TEACHER = False +_SC_HALVE = False # bipolar sign-magnitude halving during calibration probes def _save_fp_weights(model) -> None: @@ -337,6 +338,12 @@ def __init__( # group's pre-subsample row count R_g, used for the R_g-weighted global lambda. self.budget_scope = budget_scope self.true_counts: dict[tuple[str, int, int], int] = defaultdict(int) + # Per-operator MACs-per-row (out_features x in_features for linears, + # N x head_dim for qk/av). The GLOBAL budget prices each row by + # macs_per_row, so the budget is FLOP-weighted rather than row-weighted: + # a target avg stoc_len then means that value in real compute, not per + # row. Set once per operator by add(); constant across its rows. + self.op_macs: dict[str, float] = {} self._rng = np.random.default_rng(rng_seed) # Optional full-resolution log keyed by raw (operator, block_idx, timestep) so the # bucketing can be re-derived offline (any timestep/layer_buckets) without re-probing. @@ -365,10 +372,15 @@ def add( timestep: int, metric_norm: torch.Tensor, errors_by_level: list[torch.Tensor], + macs_per_row: float = 1.0, ): if not self.should_record(operator, block_idx, timestep): return + # Record this operator's per-row MAC cost (constant across its rows) so + # the global budget can weight rows by FLOPs, not just row count. + self.op_macs[operator] = float(macs_per_row) + metrics = metric_norm.detach().float().reshape(-1).cpu().numpy() err = torch.stack([e.detach().float().reshape(-1) for e in errors_by_level], dim=-1) errors = err.cpu().numpy() @@ -475,14 +487,20 @@ def export(self) -> tuple[dict, list[dict]]: # independently pinned to budget_ratio*ref -- original behaviour). global_lam = None if self.budget_scope == "global": - werrs, row_counts = [], [] + werrs, flop_counts = [], [] for key, rec in self.records.items(): e = np.concatenate(rec["errors"], axis=0) werrs.append(e) - row_counts.append(float(self.true_counts.get(key, e.shape[0]))) + R = float(self.true_counts.get(key, e.shape[0])) + # FLOP-weighted budget: price each true row by macs_per_row so + # cheap-but-numerous ops (per-head-per-row av/qk) can't dominate + # the budget over few-but-expensive linear rows. _global_lambda + # treats the passed counts as the budget/cost weights, so passing + # R*macs makes the realized average FLOP-weighted. + flop_counts.append(R * self.op_macs.get(key[0], 1.0)) global_lam = _global_lambda( werrs, self.costs, self.budget_ratio, - float(self.budget_ref_stoc_len), row_counts) + float(self.budget_ref_stoc_len), flop_counts) payload["global_lambda"] = float(global_lam) for operator in sorted(self.operators): @@ -497,7 +515,8 @@ def export(self) -> tuple[dict, list[dict]]: fitted = self._fit_group( np.concatenate(op_metrics, axis=0), np.concatenate(op_errors, axis=0), - lam=global_lam, weight=1.0, rep=self._rep_for_op(operator), + lam=global_lam, weight=1.0, + rep=self._rep_for_op(operator) * self.op_macs.get(operator, 1.0), ) payload["operator_defaults"][operator] = fitted summary_rows.append( @@ -518,8 +537,9 @@ def export(self) -> tuple[dict, list[dict]]: if metrics.size < self.min_bucket_units: continue operator, t_bucket, l_bucket = key - fitted = self._fit_group(metrics, errors, lam=global_lam, - weight=1.0, rep=self._rep_for_key(key)) + fitted = self._fit_group( + metrics, errors, lam=global_lam, weight=1.0, + rep=self._rep_for_key(key) * self.op_macs.get(key[0], 1.0)) bucket_key = f"{operator}:t{t_bucket}:l{l_bucket}" payload["buckets"][bucket_key] = fitted summary_rows.append( @@ -532,16 +552,19 @@ def export(self) -> tuple[dict, list[dict]]: } ) - # Honest iso-budget check: row-weighted (R_g) average stoc_len across all - # buckets. For per_bucket this is ~budget_ratio*ref by construction; for - # global it's the realized average -- should also land near the target. + # Honest iso-budget check: FLOP-weighted (R_g * macs_g) average stoc_len + # across all buckets. For global scope this is the realized compute + # average -- should land near the target (unlike the old row-weighted + # average, which under-counted the FLOP-heavy linears). + payload["op_macs"] = dict(self.op_macs) num = den = 0.0 for bkey, fitted in payload["buckets"].items(): op, tpart, lpart = bkey.split(":") k = (op, int(tpart[1:]), int(lpart[1:])) R = float(self.true_counts.get(k, fitted["num_units"])) - num += R * fitted["avg_stoc_len"] - den += R + w = R * self.op_macs.get(op, 1.0) + num += w * fitted["avg_stoc_len"] + den += w payload["expected_avg_stoc_len"] = (num / den) if den else 0.0 return payload, summary_rows @@ -605,6 +628,7 @@ def _run_attention_linear_level( chunk_d=chunk_d, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) result = chunk_result if result is None else result + chunk_result else: @@ -620,6 +644,7 @@ def _run_attention_linear_level( group_b=1, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) if bias is not None: @@ -674,6 +699,7 @@ def _run_mlp_linear_level( chunk_d=chunk_d, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) result = chunk_result if result is None else result + chunk_result else: @@ -690,6 +716,7 @@ def _run_mlp_linear_level( chunk_d=chunk_d, stoc_len=stoc_len, rng_levels=rng_levels, + halve_bipolar_stoc_len=_SC_HALVE, ) if bias is not None: result = result + bias @@ -758,7 +785,9 @@ def _attn_hook(self, module, inputs, _output): teacher_qkv.float().reshape(-1, teacher_qkv.shape[-1]), ) ) - self.calibrator.add("input_proj", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "input_proj", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.qkv.weight.shape[0] * module.qkv.weight.shape[1])) else: if _USE_FP_TEACHER: x_for_teacher = (torch.index_select(x, 2, module.reorder_index_qkv) @@ -784,16 +813,32 @@ def _attn_hook(self, module, inputs, _output): teacher_attn = (q_scaled @ k.transpose(-2, -1)) if self.calibrator.should_record("qk", block_idx, timestep): - q_metric = _normalize_metric(q_scaled.float().abs().amax(dim=(0, 2, 3))) - level_errors = [] + # Per-query-row recording — matches the runtime per-row qk MP + # (rows classified by ||Q_row||_inf), mirroring the av path. + q_flat = q_scaled.reshape(-1, q_scaled.shape[-2], q_scaled.shape[-1]) + teacher_flat = teacher_attn.reshape( + -1, teacher_attn.shape[-2], teacher_attn.shape[-1]) + qk_errors_by_level: list[list[torch.Tensor]] = [[] for _ in self.calibrator.levels] for sl in self.calibrator.levels: if sl == 0: sc_attn = torch.zeros_like(teacher_attn, dtype=torch.float32) else: sc_prec = _resolve_level_sc_prec(module, sl) sc_attn = module._sc_qk_uniform(q, k, sc_prec, sl).float() - level_errors.append(_METRIC_HEADS(sc_attn, teacher_attn.float())) - self.calibrator.add("qk", block_idx, timestep, q_metric, level_errors) + sc_flat = sc_attn.reshape(-1, sc_attn.shape[-2], sc_attn.shape[-1]) + for i in range(sc_flat.shape[0]): + qk_errors_by_level[self.calibrator.levels.index(sl)].append( + _METRIC_ROWS(sc_flat[i], teacher_flat[i].float()) + ) + qk_macs = float(n_tokens * module.head_dim) + for i in range(q_flat.shape[0]): + row_metric = _normalize_metric(q_flat[i].abs().amax(dim=-1)) + level_errors = [ + qk_errors_by_level[level_idx][i] + for level_idx in range(len(self.calibrator.levels)) + ] + self.calibrator.add("qk", block_idx, timestep, row_metric, + level_errors, macs_per_row=qk_macs) attn = module.attn_drop(teacher_attn.softmax(dim=-1)) teacher_av = (attn @ v) @@ -820,7 +865,9 @@ def _attn_hook(self, module, inputs, _output): level_errors_by_level[level_idx][i] for level_idx in range(len(self.calibrator.levels)) ] - self.calibrator.add("av", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "av", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.head_dim * n_tokens)) proj_hidden = module.num_heads * module.head_dim proj_in = teacher_av.transpose(1, 2).reshape(bsz, n_tokens, proj_hidden) @@ -850,7 +897,9 @@ def _attn_hook(self, module, inputs, _output): teacher_proj.float().reshape(-1, teacher_proj.shape[-1]), ) ) - self.calibrator.add("proj", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "proj", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.proj.weight.shape[0] * module.proj.weight.shape[1])) def _mlp_hook(self, module, inputs, _output): timestep = module.sc_controller.current_timestep @@ -887,7 +936,9 @@ def _mlp_hook(self, module, inputs, _output): teacher_fc1.float().reshape(-1, teacher_fc1.shape[-1]), ) ) - self.calibrator.add("mlp_fc1", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "mlp_fc1", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.fc1.weight.shape[0] * module.fc1.weight.shape[1])) else: if _USE_FP_TEACHER: x_for_teacher = (torch.index_select(x, 2, module.reorder_index_fc1) @@ -919,7 +970,9 @@ def _mlp_hook(self, module, inputs, _output): teacher_fc2.float().reshape(-1, teacher_fc2.shape[-1]), ) ) - self.calibrator.add("mlp_fc2", block_idx, timestep, row_metric, level_errors) + self.calibrator.add( + "mlp_fc2", block_idx, timestep, row_metric, level_errors, + macs_per_row=float(module.fc2.weight.shape[0] * module.fc2.weight.shape[1])) def register_hooks(self): hooks = [] @@ -1111,11 +1164,13 @@ def main(): model = quantize_sc_model(model, device, args, sc_controller=sc_controller) # Wire metric + teacher dispatch globals. - global _METRIC_ROWS, _METRIC_HEADS, _USE_FP_TEACHER + global _METRIC_ROWS, _METRIC_HEADS, _USE_FP_TEACHER, _SC_HALVE if args.metric == "cosine": _METRIC_ROWS = _cosine_dist_rows _METRIC_HEADS = _cosine_dist_heads _USE_FP_TEACHER = (args.teacher == "fp") + _SC_HALVE = bool(args.sc_halve) + print(f"SC halve during calibration: {_SC_HALVE}") print(f"Calibration metric: {args.metric}, teacher: {args.teacher}") model.to(device) model.eval().half() diff --git a/scripts/quant_sc_main.py b/scripts/quant_sc_main.py index 7e07f3d..9159685 100644 --- a/scripts/quant_sc_main.py +++ b/scripts/quant_sc_main.py @@ -424,20 +424,12 @@ def main(): quant_string_name += "_fixlvlprec" if args.sc_noise_model: quant_string_name += "_noisemodel" - # Append per-operator MP alpha/beta to folder name when adaptive_mp is used + # Tag the folder name with the calibrated MP table when adaptive_mp is used if getattr(args, 'adaptive_mp', False) or getattr(args, 'adaptive_mp_table', None): - _mp_parts = [] - for _op in ["qk", "av", "proj", "input_proj", "mlp_fc1", "mlp_fc2"]: - _a = getattr(args, f"mp_alpha_{_op}", None) - _b = getattr(args, f"mp_beta_{_op}", None) - if _a is not None and _b is not None: - _mp_parts.append(f"{_op}_a{_a}_b{_b}") if getattr(args, 'adaptive_mp_table', None): quant_string_name += f"_mptbl_{Path(args.adaptive_mp_table).stem}" - elif _mp_parts: - quant_string_name += "_mp_" + "_".join(_mp_parts) else: - quant_string_name += f"_mp_a{args.mp_alpha}_b{args.mp_beta}" + quant_string_name += "_adaptivemp" if getattr(args, 'range_mp', False): _rmp_parts = [] for _op in ["qk", "av", "proj", "input_proj", "mlp"]: @@ -545,44 +537,20 @@ def main(): # Initialize mixed precision if requested if args.adaptive_mp or args.adaptive_mp_table: + if not args.adaptive_mp_table: + raise ValueError( + "--adaptive_mp requires --adaptive_mp_table (a calibrated " + "threshold table from calibrate_mp_thresholds.py). The old " + "closed-form alpha/beta fallback has been removed.") levels = [int(x) for x in args.mp_levels.split(',')] - # Build per-operator overrides (None means use global default) - operator_params = {} - # Lookup order: exact op name → group key → global - for op, group_key in [("qk", "qk"), ("av", "av"), - ("mlp_fc1", "mlp_fc1"), ("mlp_fc2", "mlp_fc2"), - ("input_proj", "input_proj"), ("proj", "proj")]: - # Try exact op-level arg first, then group-level fallback - a = getattr(args, f"mp_alpha_{op}", None) - b = getattr(args, f"mp_beta_{op}", None) - if a is None: - # Group fallback: mlp_fc1/mlp_fc2 → mlp - group_fallback = {"mlp_fc1": "mlp", "mlp_fc2": "mlp"}.get(op) - if group_fallback: - a = getattr(args, f"mp_alpha_{group_fallback}", None) - if b is None: - group_fallback = {"mlp_fc1": "mlp", "mlp_fc2": "mlp"}.get(op) - if group_fallback: - b = getattr(args, f"mp_beta_{group_fallback}", None) - if a is not None or b is not None: - operator_params[op] = ( - a if a is not None else args.mp_alpha, - b if b is not None else args.mp_beta, - ) - adaptive_config = AdaptiveMPConfig( stoc_len_levels=levels, - alpha=args.mp_alpha, - beta=args.mp_beta, enable_pruning=args.mp_enable_pruning, - operator_params=operator_params, threshold_table_path=args.adaptive_mp_table, ) sc_controller.init_adaptive_mp(adaptive_config) - logging.info(f"Adaptive mixed precision V2 enabled: levels={levels}, " - f"alpha={args.mp_alpha}, beta={args.mp_beta}, " + logging.info(f"Adaptive mixed precision enabled: levels={levels}, " f"pruning={args.mp_enable_pruning}, " - f"operator_params={operator_params}, " f"threshold_table={args.adaptive_mp_table}") elif args.mp: levels = [int(x) for x in args.mp_levels.split(',')] @@ -744,25 +712,17 @@ def create_argparser(): help='Comma-separated fractions per level (default: equal). E.g. "0.1,0.2,0.3,0.4".' ) - # Adaptive mixed precision (timestep-aware, inspired by APT) + # Adaptive mixed precision (calibrated per-row thresholds) parser.add_argument( '--adaptive_mp', action='store_true', - help='Enable adaptive mixed precision with timestep-aware thresholds.' + help='Enable adaptive mixed precision. Requires --adaptive_mp_table ' + '(a calibrated threshold table); there is no closed-form fallback.' ) parser.add_argument( '--adaptive_mp_table', type=str, default=None, - help='Path to a calibrated adaptive-MP threshold JSON table. ' - 'When provided, runtime classification uses operator/timestep/layer ' - 'bucket thresholds from the table and falls back to alpha/beta only ' - 'for operators missing from the table.' - ) - parser.add_argument( - '--mp_alpha', type=float, default=0.3, - help='Adaptive MP: threshold sensitivity to timestep progress.' - ) - parser.add_argument( - '--mp_beta', type=float, default=0.05, - help='Adaptive MP: base threshold offset.' + help='Path to a calibrated adaptive-MP threshold JSON table from ' + 'calibrate_mp_thresholds.py. Runtime classification uses the ' + 'per-(operator, timestep, layer) bucket thresholds from the table.' ) parser.add_argument( '--mp_enable_pruning', action='store_true', default=True, @@ -772,64 +732,6 @@ def create_argparser(): '--no_mp_pruning', dest='mp_enable_pruning', action='store_false', help='Adaptive MP: disable row pruning.' ) - # Per-operator alpha/beta overrides - parser.add_argument( - '--mp_alpha_qk', type=float, default=None, - help='Per-operator alpha for QK (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_av', type=float, default=None, - help='Per-operator alpha for AV (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_mlp', type=float, default=None, - help='Per-operator alpha for MLP fc1/fc2 (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_proj', type=float, default=None, - help='Per-operator alpha for proj (default: use --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_input_proj', type=float, default=None, - help='Per-operator alpha for input_proj (default: fallback to --mp_alpha_proj, then --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_mlp_fc1', type=float, default=None, - help='Per-operator alpha for mlp_fc1 (default: fallback to --mp_alpha_mlp, then --mp_alpha).' - ) - parser.add_argument( - '--mp_alpha_mlp_fc2', type=float, default=None, - help='Per-operator alpha for mlp_fc2 (default: fallback to --mp_alpha_mlp, then --mp_alpha).' - ) - parser.add_argument( - '--mp_beta_qk', type=float, default=None, - help='Per-operator beta for QK (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_av', type=float, default=None, - help='Per-operator beta for AV (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_mlp', type=float, default=None, - help='Per-operator beta for MLP fc1/fc2 (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_proj', type=float, default=None, - help='Per-operator beta for proj (default: use --mp_beta).' - ) - parser.add_argument( - '--mp_beta_input_proj', type=float, default=None, - help='Per-operator beta for input_proj (default: fallback to --mp_beta_proj, then --mp_beta).' - ) - parser.add_argument( - '--mp_beta_mlp_fc1', type=float, default=None, - help='Per-operator beta for mlp_fc1 (default: fallback to --mp_beta_mlp, then --mp_beta).' - ) - parser.add_argument( - '--mp_beta_mlp_fc2', type=float, default=None, - help='Per-operator beta for mlp_fc2 (default: fallback to --mp_beta_mlp, then --mp_beta).' - ) - # Range-based mixed precision (weight min/max range) parser.add_argument( '--range_mp', action='store_true', From 7a326953ba9f3155c304beeefaea4143c3fb6713 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:55:25 -0400 Subject: [PATCH 2/6] chore: gitignore SLURM job logs and generated figures Run artifacts (slurm-*.out, figures/) are regenerated, not source. Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index c6277fb..8120ce8 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,7 @@ logs_mp_sweep/ evaluation/sample_grids/ evaluation/imagenet_ref/images/ evaluation/imagenet_ref/*.npz + +# SLURM job logs and generated figures (run artifacts — regenerated, not source) +slurm-*.out +figures/ From 0b66c15d0897b4da71b82a8319619105d7ef64fc Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:55:47 -0400 Subject: [PATCH 3/6] tools+scripts: add SC comparison tools and MP/calib sbatch scripts tools/: compare_old_vs_new{,_isolated,_thorough}.py, compare_vs_cpu_reference.py, sc_call_counter{,_full}.py, smoke_test_e2e.py scripts/: global-PR calib + MP sbatch (calib_globalpr{,_flop}{,_halve}, mp_globalpr_{auto,fill,flop_halve_auto}, uniform{,_halve}_auto, opfreeze_probe) + eval/ CPU eval Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/eval/eval_mp_flop_halve_cpu.sh | 37 +++ scripts/eval/sbatch_eval_mp_flop_halve.sb | 42 ++++ scripts/sbatch_calib_globalpr.sb | 31 +++ scripts/sbatch_calib_globalpr_flop.sb | 32 +++ scripts/sbatch_calib_globalpr_flop_halve.sb | 33 +++ scripts/sbatch_calib_globalpr_halve.sb | 31 +++ scripts/sbatch_mp_globalpr_auto.sb | 56 +++++ scripts/sbatch_mp_globalpr_fill.sb | 56 +++++ scripts/sbatch_mp_globalpr_flop_halve_auto.sb | 56 +++++ scripts/sbatch_opfreeze_probe.sb | 45 ++++ scripts/sbatch_uniform_auto.sb | 49 ++++ scripts/sbatch_uniform_halve_auto.sb | 49 ++++ tools/compare_old_vs_new.py | 207 ++++++++++++++++ tools/compare_old_vs_new_isolated.py | 119 ++++++++++ tools/compare_old_vs_new_thorough.py | 167 +++++++++++++ tools/compare_vs_cpu_reference.py | 156 ++++++++++++ tools/sc_call_counter.py | 61 +++++ tools/sc_call_counter_full.py | 59 +++++ tools/smoke_test_e2e.py | 223 ++++++++++++++++++ 19 files changed, 1509 insertions(+) create mode 100644 scripts/eval/eval_mp_flop_halve_cpu.sh create mode 100644 scripts/eval/sbatch_eval_mp_flop_halve.sb create mode 100644 scripts/sbatch_calib_globalpr.sb create mode 100644 scripts/sbatch_calib_globalpr_flop.sb create mode 100644 scripts/sbatch_calib_globalpr_flop_halve.sb create mode 100755 scripts/sbatch_calib_globalpr_halve.sb create mode 100644 scripts/sbatch_mp_globalpr_auto.sb create mode 100755 scripts/sbatch_mp_globalpr_fill.sb create mode 100644 scripts/sbatch_mp_globalpr_flop_halve_auto.sb create mode 100644 scripts/sbatch_opfreeze_probe.sb create mode 100644 scripts/sbatch_uniform_auto.sb create mode 100755 scripts/sbatch_uniform_halve_auto.sb create mode 100644 tools/compare_old_vs_new.py create mode 100644 tools/compare_old_vs_new_isolated.py create mode 100644 tools/compare_old_vs_new_thorough.py create mode 100644 tools/compare_vs_cpu_reference.py create mode 100644 tools/sc_call_counter.py create mode 100644 tools/sc_call_counter_full.py create mode 100644 tools/smoke_test_e2e.py diff --git a/scripts/eval/eval_mp_flop_halve_cpu.sh b/scripts/eval/eval_mp_flop_halve_cpu.sh new file mode 100644 index 0000000..0135f64 --- /dev/null +++ b/scripts/eval/eval_mp_flop_halve_cpu.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# CPU-only eval for one FLOP+halve adaptive-MP config (no GPU -> respects the +# 6-GPU cap). KID + FID/IS/Precision/Recall vs the ImageNet-256 reference. +# Usage: bash eval_mp_flop_halve_cpu.sh +set -uo pipefail +AVG="${1:?usage: eval_mp_flop_halve_cpu.sh }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +PACKER=$PREV/imagenet256_ref/parallel_npz.py +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +BASE=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +AD=$BASE/samples +WORK=$BASE/_eval; mkdir -p "$WORK" +OUT=$WORK/eval_avg${AVG}.txt +export CUDA_VISIBLE_DEVICES="" # force CPU +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$OUT") 2>&1 +echo "############ FLOP+HALVE MP EVAL avg${AVG} (CPU) $(date) ############" + +IDX=$WORK/idx.txt; NPZ=$WORK/avg${AVG}.npz +(cd "$AD" && ls [0-9][0-9][0-9][0-9][0-9][0-9].png | awk '($1+0)<2000' | sort) > "$IDX" +SEL=$WORK/sel; rm -rf "$SEL"; mkdir -p "$SEL" +while read -r f; do [[ -n "$f" ]] && ln -sf "$AD/$f" "$SEL/$f"; done < "$IDX" +echo "=== matched idx 0-1999: $(wc -l < "$IDX") ===" + +conda activate qdit +python -u "$PACKER" "$SEL" "$NPZ" +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +echo ""; echo "===== [1] KID vs ref (idx 0-1999) =====" +bash "$REPO/scripts/eval/kid_openai.sh" "$WORK/kid_avg${AVG}.txt" "$NPZ" +echo ""; echo "===== [2] FID / IS / sFID / Precision / Recall vs ref =====" +python -u "$EVALUATOR" "$REF" "$NPZ" 2>&1 | grep -E "^(Inception Score|FID|sFID|Precision|Recall):" | sed "s/^/[avg${AVG}] /" +conda deactivate +echo ""; echo "############ DONE avg${AVG} $(date) ############" diff --git a/scripts/eval/sbatch_eval_mp_flop_halve.sb b/scripts/eval/sbatch_eval_mp_flop_halve.sb new file mode 100644 index 0000000..0a4cfa0 --- /dev/null +++ b/scripts/eval/sbatch_eval_mp_flop_halve.sb @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=eval_mp +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:45:00 +# GPU eval for one FLOP+halve adaptive-MP config: KID + FID/IS/sFID/Prec/Recall. +# Usage: sbatch sbatch_eval_mp_flop_halve.sb +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_eval_mp_flop_halve.sb }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +PACKER=$PREV/imagenet256_ref/parallel_npz.py +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +BASE=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +AD=$BASE/samples +WORK=$BASE/_eval; mkdir -p "$WORK" +OUT=$WORK/eval_avg${AVG}.txt +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$OUT") 2>&1 +echo "############ FLOP+HALVE MP EVAL avg${AVG} (GPU) $(date) ############" + +IDX=$WORK/idx.txt; NPZ=$WORK/avg${AVG}.npz +(cd "$AD" && ls [0-9][0-9][0-9][0-9][0-9][0-9].png | awk '($1+0)<2000' | sort) > "$IDX" +SEL=$WORK/sel; rm -rf "$SEL"; mkdir -p "$SEL" +while read -r f; do [[ -n "$f" ]] && ln -sf "$AD/$f" "$SEL/$f"; done < "$IDX" +echo "=== matched idx 0-1999: $(wc -l < "$IDX") ===" + +conda activate qdit +python -u "$PACKER" "$SEL" "$NPZ" +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +echo ""; echo "===== [1] KID vs ref (idx 0-1999) =====" +bash "$REPO/scripts/eval/kid_openai.sh" "$WORK/kid_avg${AVG}.txt" "$NPZ" +echo ""; echo "===== [2] FID / IS / sFID / Precision / Recall vs ref =====" +python -u "$EVALUATOR" "$REF" "$NPZ" 2>&1 | grep -E "^(Inception Score|FID|sFID|Precision|Recall):" | sed "s/^/[avg${AVG}] /" +conda deactivate +echo ""; echo "############ DONE avg${AVG} $(date) ############" diff --git a/scripts/sbatch_calib_globalpr.sb b/scripts/sbatch_calib_globalpr.sb new file mode 100644 index 0000000..f32bc9c --- /dev/null +++ b/scripts/sbatch_calib_globalpr.sb @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=calib_gpr +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# per-row GLOBAL calibration for one budget. Usage: sbatch sbatch_calib_globalpr.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== calib GLOBAL per-row avg${AVG} br=${BR} $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref192.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_flop.sb b/scripts/sbatch_calib_globalpr_flop.sb new file mode 100644 index 0000000..6e5113a --- /dev/null +++ b/scripts/sbatch_calib_globalpr_flop.sb @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=calib_flop +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# FLOP-weighted per-row GLOBAL calibration for one budget (fix branch: +# fix/flop-cost-perrow-qk-drop-alphabeta). Usage: sbatch sbatch_calib_globalpr_flop.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== FLOP-weighted calib GLOBAL per-row avg${AVG} br=${BR} branch=$(git branch --show-current) $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref256.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_flop_halve.sb b/scripts/sbatch_calib_globalpr_flop_halve.sb new file mode 100644 index 0000000..06d1984 --- /dev/null +++ b/scripts/sbatch_calib_globalpr_flop_halve.sb @@ -0,0 +1,33 @@ +#!/bin/bash +#SBATCH --job-name=calib_flopH +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# FLOP-weighted + HALVE per-row GLOBAL calibration for one budget (fix branch: +# fix/flop-cost-perrow-qk-drop-alphabeta). Matches the halve MP deployment. +# Usage: sbatch sbatch_calib_globalpr_flop_halve.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== FLOP+HALVE calib GLOBAL per-row avg${AVG} br=${BR} branch=$(git branch --show-current) $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_halve.sb b/scripts/sbatch_calib_globalpr_halve.sb new file mode 100755 index 0000000..d2f0403 --- /dev/null +++ b/scripts/sbatch_calib_globalpr_halve.sb @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=calib_gpr +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=8 +#SBATCH --mem=110G +#SBATCH --time=03:00:00 +# per-row GLOBAL calibration for one budget. Usage: sbatch sbatch_calib_globalpr.sb +set -uo pipefail +AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_halve +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" +echo "=== calib GLOBAL per-row HALVE avg${AVG} br=${BR} $(date) ===" +python -u scripts/calibrate_mp_thresholds.py \ + --mp_levels 256,192,128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --metric cosine --teacher fp \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --image-size 256 --num-sampling-steps 50 \ + --num_calib_batches 1 --num_calib_timesteps 50 \ + --timestep_buckets 10 --layer_buckets 4 \ + --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref192.json" \ + --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" +echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_mp_globalpr_auto.sb b/scripts/sbatch_mp_globalpr_auto.sb new file mode 100644 index 0000000..e358c2a --- /dev/null +++ b/scripts/sbatch_mp_globalpr_auto.sb @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=mp_gpr_auto +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# +# global+per-row generation that AUTO-RESUMES to NUM_FID=10000, TIMEOUT-safe: +# it queues its successor (afterany, chain-capped) at the START, before it can +# time out. Usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow/calib_fix_avg${AVG}_l256_ref192.json +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_cfg15/adaptive_avg${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== avg${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete] avg${AVG} already at $NUM_FID"; exit 0; fi +# Queue successor NOW (afterany -> runs after me regardless of timeout/complete); chain-capped. +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + && echo "[chain] queued successor (chain=$((CHAIN+1)))" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== avg${AVG}: $START -> $(cnt) / $NUM_FID at $(date) ===" diff --git a/scripts/sbatch_mp_globalpr_fill.sb b/scripts/sbatch_mp_globalpr_fill.sb new file mode 100755 index 0000000..b88939a --- /dev/null +++ b/scripts/sbatch_mp_globalpr_fill.sb @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=mp_gpr_fill +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# +# global+per-row generation that AUTO-RESUMES to NUM_FID=10000, TIMEOUT-safe: +# it queues its successor (afterany, chain-capped) at the START, before it can +# time out. Usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_mp_globalpr_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_fill.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow/calib_fix_avg${AVG}_l256_ref192.json +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_cfg15/adaptive_avg${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' -printf '%f\n' | awk -v N=$NUM_FID '{n=$1+0; if(n runs after me regardless of timeout/complete); chain-capped. +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + && echo "[chain] queued successor (chain=$((CHAIN+1)))" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== avg${AVG}: $START -> $(cnt) / $NUM_FID at $(date) ===" diff --git a/scripts/sbatch_mp_globalpr_flop_halve_auto.sb b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb new file mode 100644 index 0000000..7debbc4 --- /dev/null +++ b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=mp_flopH +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# +# FLOP-weighted + HALVE adaptive-MP generation, auto-resumes to NUM_FID, TIMEOUT-safe: +# queues its successor (afterany, chain-capped) at the START. Uses the FLOP+halve +# calibrated tables. Usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_flop_halve_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg${AVG}.json +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== FLOP+HALVE avg${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete] avg${AVG} already at $NUM_FID"; exit 0; fi +# Queue successor NOW (afterany -> runs after me regardless of timeout/complete); chain-capped. +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + && echo "[chain] queued successor (chain=$((CHAIN+1)))" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== FLOP+HALVE avg${AVG}: $START -> $(cnt) / $NUM_FID at $(date) ===" diff --git a/scripts/sbatch_opfreeze_probe.sb b/scripts/sbatch_opfreeze_probe.sb new file mode 100644 index 0000000..3342bfb --- /dev/null +++ b/scripts/sbatch_opfreeze_probe.sb @@ -0,0 +1,45 @@ +#!/bin/bash +#SBATCH --job-name=opfreeze +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem-per-gpu=16G +#SBATCH --time=04:00:00 +# +# Op-freeze sensitivity probe: all ops at the avg48 FLOP-halve budget EXCEPT one +# op forced to stoc_len 256 (frozen-to-FP-precision table). Generates the same +# indices 0-63 (seed-matched to FP & the avg48 baseline) so PSNR-vs-FP is paired. +# Usage: sbatch sbatch_opfreeze_probe.sb (OP in qk av input_proj proj mlp_fc1 mlp_fc2) +set -uo pipefail +OP="${1:?usage: sbatch sbatch_opfreeze_probe.sb }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_opfreeze_avg48/freeze_${OP}.json +PROBE=$SCRATCH/scmp_diffusion_opfreeze_avg48 +OUT=$PROBE/freeze_${OP}; SAMPLES=$OUT/samples; LOG=$OUT/_logs +IDXF="${2:-$PROBE/idx_0_23.txt}" +N=$(grep -c . "$IDXF") +BALANCED=10000; BATCH=$N; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=256,192,128,96,64,48,32,16 +mkdir -p "$SAMPLES" "$LOG" +[[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } +[[ -s "$IDXF" ]] || { echo "ERROR: missing index file $IDXF" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=4; cd "$REPO" +have(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +echo "=== opfreeze OP=$OP start=$(have)/$N job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$(have)" -ge "$N" ]]; then echo "[complete] $OP already at $N"; exit 0; fi +CUDA_VISIBLE_DEVICES=0 python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row --sc_halve \ + --adaptive_mp --adaptive_mp_table "$CALIB" --mp_levels "$MP_LEVELS" \ + --image-size 256 --num-sampling-steps "$STEPS" --cfg-scale "$CFG" --batch-size "$BATCH" \ + --generate-fid-samples --balanced_classes --num-classes "$NUM_CLASSES" \ + --balanced_total_samples "$BALANCED" --num-fid-samples "$BALANCED" \ + --target_indices_path "$IDXF" --samples_dir_override "$SAMPLES" \ + --seed "$SEED" --results-dir "$LOG" --ckpt "$CKPT" \ + 2>&1 | tee "$LOG/run.log" +echo "=== opfreeze OP=$OP done=$(have)/$N $(date) ===" diff --git a/scripts/sbatch_uniform_auto.sb b/scripts/sbatch_uniform_auto.sb new file mode 100644 index 0000000..77e1ab5 --- /dev/null +++ b/scripts/sbatch_uniform_auto.sb @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --job-name=uni_auto +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# uniform generation, TIMEOUT-safe auto-resume to 10000. Usage: sbatch sbatch_uniform_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?AVG}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_uniform_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_uniform${AVG}_all.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_bitrev_cfg15/uniform${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$SC_JSON" ]] || { echo "ERROR: missing $SC_JSON" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== uniform${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete]"; exit 0; fi +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) && echo "[chain] queued successor" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== uniform${AVG}: $START -> $(cnt)/$NUM_FID $(date) ===" diff --git a/scripts/sbatch_uniform_halve_auto.sb b/scripts/sbatch_uniform_halve_auto.sb new file mode 100755 index 0000000..f3fa9eb --- /dev/null +++ b/scripts/sbatch_uniform_halve_auto.sb @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --job-name=uni_halve +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=1-12:00:00 +# uniform generation, TIMEOUT-safe auto-resume to 10000. Usage: sbatch sbatch_uniform_auto.sb [CHAIN] +set -uo pipefail +AVG="${1:?AVG}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_uniform_halve_auto.sb +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_uniform${AVG}_all.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_halveON_cfg15/uniform${AVG} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$SC_JSON" ]] || { echo "ERROR: missing $SC_JSON" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +cnt(){ find "$SAMPLES" -maxdepth 1 -name '[0-9][0-9][0-9][0-9][0-9][0-9].png' | wc -l; } +START=$(cnt) +echo "=== uniform${AVG} auto: start=$START/$NUM_FID chain=$CHAIN job=${SLURM_JOB_ID:-?} $(date) ===" +if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete]"; exit 0; fi +if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then + sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) && echo "[chain] queued successor" +fi +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== uniform${AVG}: $START -> $(cnt)/$NUM_FID $(date) ===" diff --git a/tools/compare_old_vs_new.py b/tools/compare_old_vs_new.py new file mode 100644 index 0000000..fe5483d --- /dev/null +++ b/tools/compare_old_vs_new.py @@ -0,0 +1,207 @@ +"""Numerical comparison: scmp_llm/SC/sc_triton.py (old) vs scmp_kernels.sc.sc_matmul (new). + +For each shape pattern and mode, runs both implementations on identical inputs +with identical Sobol config + identical stoc_len, then reports: + - max abs diff between the two SC outputs + - max rel diff between the two SC outputs + - rel_err of each vs torch.matmul (the fp baseline) + +Expected outcome: outputs differ slightly due to the clipping margin removal +(bipolar ±125→±127, unipolar [2,253]→[0,255]). Both should be within sane +SC noise distance from the fp baseline. +""" +from __future__ import annotations +import sys, os +from pathlib import Path + +# Old impl — bare imports relative to SC/ folder +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_triton as old_sc +from config_helpers import make_sobol_simple_config + +# New impl — installed as scmp_kernels package +from scmp_kernels.sc import sc_matmul as new_sc_matmul, clear_rng_cache +import scmp_kernels.sc.kernels as new_sc_kernels + +import torch + + +DEVICE = "cuda" + + +def diff(a, b): + """Return (max_abs, max_rel).""" + d = (a - b).abs() + rel = d / b.abs().clamp_min(1e-6) + return float(d.max().item()), float(rel.max().item()) + + +def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + +def line(label, vals): + print(f" {label:<60} " + " ".join(f"{v:>10}" for v in vals), flush=True) + + +def header(title): + print(f"\n{title}\n{'='*120}", flush=True) + + +def reset(): + old_sc.clear_rng_cache() + clear_rng_cache() + + +def run_case(name, a, b, *, mode, granularity_new, old_call, sc_prec=8, stoc_len=256): + """Run both impls on (a, b), print diff vs each other + each vs fp.""" + D = a.shape[-1] + cfg = make_sobol_simple_config(D, D, sc_prec) + fp = a @ b.transpose(-1, -2) if b.dim() == a.dim() else a @ b.t() + + reset() + out_old = old_call(a, b, cfg, sc_prec, stoc_len, mode) + + reset() + out_new = new_sc_matmul(a, b, granularity=granularity_new, mode=mode, + sc_prec=sc_prec, stoc_len=stoc_len, config=cfg) + + assert out_old.shape == out_new.shape == fp.shape, \ + f"{name}: shape mismatch old={out_old.shape} new={out_new.shape} fp={fp.shape}" + + max_abs, max_rel = diff(out_new, out_old) + re_old = rel_err(out_old, fp) + re_new = rel_err(out_new, fp) + bitwise = "BIT-IDENTICAL" if max_abs == 0 else ("≈identical" if max_abs < 1e-5 else "differ") + line(name, [f"{max_abs:.3e}", f"{max_rel:.3e}", + f"{re_old:.4f}", f"{re_new:.4f}", bitwise]) + + +# -------------------------------------------------------------------------- +# Adapters — each wraps an old-impl call to a uniform signature +# -------------------------------------------------------------------------- +def _old_enable_triton(a, b, cfg, sc_prec, stoc_len, mode): + return old_sc.sc_matmul_enable_triton( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + + +def _old_enable_batched_bipolar(q, k, cfg, sc_prec, stoc_len, mode): + q_maxs = q.amax(dim=(1, 2)) + q_mins = q.amin(dim=(1, 2)) + k_maxs = k.amax(dim=(1, 2)) + k_mins = k.amin(dim=(1, 2)) + return old_sc.sc_matmul_enable_batched_bipolar( + q, k, q_maxs, q_mins, k_maxs, k_mins, + sc_prec, cfg, stoc_len=stoc_len, + ) + + +def _old_grouped_enable(a, b, cfg, sc_prec, stoc_len, mode): + return old_sc.sc_matmul_grouped_enable_triton( + a, b, group_a=a.shape[0], group_b=b.shape[0], + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + + +def main(): + print(f"scmp_llm SC source: {SCMP_LLM_SC}", flush=True) + print(f"new scmp_kernels: {Path(new_sc_kernels.__file__).parent}", flush=True) + print(f"device: {torch.cuda.get_device_name(0)}", flush=True) + + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + torch.manual_seed(0) + + # 2D matmul: per_tensor bipolar — sc_matmul_enable_triton vs sc_matmul(granularity="per_tensor") + header("BIPOLAR — 2D matmul A(N×D) @ B(M×D).T ─ enable-signal table-lookup") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for shape in [(32, 64, 16), (32, 128, 64), (128, 1152, 1152)]: + N, D, M = shape + a = torch.randn(N, D, device=DEVICE) + b = torch.randn(M, D, device=DEVICE) * 0.1 + run_case(f"matmul N={N:<4} D={D:<5} M={M:<5} bipolar per_tensor", + a, b, mode="bipolar", granularity_new="per_tensor", + old_call=_old_enable_triton) + + # 2D unipolar + header("UNIPOLAR — 2D matmul A(N×D) @ B(M×D).T") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for shape in [(32, 64, 16), (32, 128, 64)]: + N, D, M = shape + a = torch.rand(N, D, device=DEVICE) # unipolar wants non-neg-ish + b = torch.randn(M, D, device=DEVICE) * 0.1 + run_case(f"matmul N={N:<4} D={D:<5} M={M:<5} unipolar per_tensor", + a, b, mode="unipolar", granularity_new="per_tensor", + old_call=_old_enable_triton) + + # GEMV — rank-1 a (matrix-vector): A(1×D) @ B(M×D).T → (1, M) + header("GEMV — A(1×D) @ B(M×D).T ─ same path, just N=1") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for D, M in [(64, 16), (128, 64), (1152, 1152)]: + a = torch.randn(1, D, device=DEVICE) + b = torch.randn(M, D, device=DEVICE) * 0.1 + run_case(f"GEMV D={D:<5} M={M:<5} bipolar per_tensor", + a, b, mode="bipolar", granularity_new="per_tensor", + old_call=_old_enable_triton) + + # 3D per-head bipolar — QK pattern in attention + header("3D BATCHED BIPOLAR — Q(BH×N×D) @ K(BH×N×D).T ─ per-head SC") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for BH, N, D in [(8, 32, 64), (16, 256, 72)]: + q = torch.randn(BH, N, D, device=DEVICE) + k = torch.randn(BH, N, D, device=DEVICE) + run_case(f"QK BH={BH:<3} N={N:<5} D={D:<5} bipolar per_head", + q, k, mode="bipolar", granularity_new="per_head", + old_call=_old_enable_batched_bipolar) + + # AV-style grouped matmul (per-row groups) + header("GROUPED — A(N×D) @ B(M×D).T ─ per-row groups, group_a=N, group_b=M") + line("CASE", ["max|Δ|", "max rel-Δ", "rel_err old", "rel_err new", "vs old"]) + print("-" * 120, flush=True) + + for N, D, M in [(64, 32, 128), (256, 72, 256)]: + a = torch.softmax(torch.randn(N, M, device=DEVICE), dim=-1) # softmax-like + b = torch.randn(D, M, device=DEVICE) * 0.1 + N2 = a.shape[0]; M2 = b.shape[0] + # new dispatcher: granularity=per_row + explicit group_a/group_b + cfg = make_sobol_simple_config(M, M, 8) + fp = a @ b.t() + reset() + out_old = old_sc.sc_matmul_grouped_enable_triton( + a, b, group_a=N, group_b=D, mode="bipolar", sc_prec=8, config=cfg, stoc_len=256) + reset() + out_new = new_sc_matmul( + a, b, granularity="per_row", group_a=N, group_b=D, + mode="bipolar", sc_prec=8, stoc_len=256, config=cfg) + ma, mr = diff(out_new, out_old) + re_o = rel_err(out_old, fp); re_n = rel_err(out_new, fp) + bit = "BIT-IDENTICAL" if ma == 0 else ("≈identical" if ma < 1e-5 else "differ") + line(f"grouped N={N:<4} D={D:<4} M={M:<4} bipolar per_row", + [f"{ma:.3e}", f"{mr:.3e}", f"{re_o:.4f}", f"{re_n:.4f}", bit]) + + print("\nLegend:") + print(" max|Δ| = max |out_new - out_old| (0 ⇒ bit-identical)") + print(" max rel-Δ = max |out_new - out_old| / |out_old|") + print(" rel_err = SC vs fp baseline (lower = closer to fp)") + print() + print("Note: any non-zero Δ is explained entirely by the clipping-margin removal") + print(" (bipolar: ±125 → ±127, unipolar: [2,253] → [0,255]) which uses more") + print(" quantization levels — so the new version's rel_err should be ≤ the old.") + + +if __name__ == "__main__": + main() diff --git a/tools/compare_old_vs_new_isolated.py b/tools/compare_old_vs_new_isolated.py new file mode 100644 index 0000000..e805feb --- /dev/null +++ b/tools/compare_old_vs_new_isolated.py @@ -0,0 +1,119 @@ +"""Isolate the cause of the bipolar per_tensor diff. + +Monkey-patches scmp_llm's ``fused_quantize_bipolar`` to use ``q_clip = q_norm`` +(no margin). If this is the sole source of the diff, the output must become +bit-identical to scmp_kernels. +""" +from __future__ import annotations +import sys +from pathlib import Path + +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_triton as old_sc +from config_helpers import make_sobol_simple_config + +from scmp_kernels.sc import sc_matmul as new_sc_matmul, clear_rng_cache + +import torch + + +# --- Save originals + install patched fused_quantize_bipolar --------------- +_orig_fqb = old_sc.fused_quantize_bipolar + + +def patched_fused_quantize_bipolar(fp_tensor, abs_max, sc_prec, rng_levels=None): + """scmp_llm's body, but with q_clip = q_norm (no margin) — matches scmp_kernels.""" + import triton + rows, cols = fp_tensor.shape + q_norm = 2 ** (sc_prec - 1) - 1 # 127 + q_clip = q_norm # 127 ← was q_norm - 2 + max_rng_val = old_sc._resolve_rng_levels(sc_prec, rng_levels) + abs_max = max(abs_max, 1e-5) + scale = abs_max / q_clip + inv_scale = 1.0 / scale + + boundary = torch.empty(rows, cols, dtype=torch.int16, device=fp_tensor.device) + sign = torch.empty(rows, cols, dtype=torch.int8, device=fp_tensor.device) + total = rows * cols + BLOCK = 1024 + grid = (triton.cdiv(total, BLOCK),) + old_sc.fused_quant_bipolar_kernel[grid]( + fp_tensor, boundary, sign, + inv_scale, q_clip, -q_clip, q_clip, max_rng_val, + rows, cols, BLOCK, + ) + return boundary, sign, scale + + +def diff(a, b): + d = (a - b).abs() + return float(d.max().item()), float((d / b.abs().clamp_min(1e-6)).max().item()) + + +def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + +def run_pair(name, a, b, *, patch): + """If patch=True, use the no-margin version of fused_quantize_bipolar.""" + D = a.shape[-1] + cfg = make_sobol_simple_config(D, D, 8) + fp = a @ b.t() + + if patch: + old_sc.fused_quantize_bipolar = patched_fused_quantize_bipolar + else: + old_sc.fused_quantize_bipolar = _orig_fqb + + old_sc.clear_rng_cache(); clear_rng_cache() + out_old = old_sc.sc_matmul_enable_triton( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="bipolar", sc_prec=8, config=cfg, stoc_len=256, + ) + + old_sc.clear_rng_cache(); clear_rng_cache() + out_new = new_sc_matmul(a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=256, config=cfg) + + ma, mr = diff(out_new, out_old) + re_o = rel_err(out_old, fp) + re_n = rel_err(out_new, fp) + bit = "BIT-IDENTICAL" if ma == 0.0 else "differ" + flag = "[PATCHED]" if patch else "[ORIGINAL]" + print(f" {flag:<11} {name:<46} max|Δ|={ma:.3e} rel_err: old={re_o:.4f} new={re_n:.4f} {bit}", flush=True) + + +def main(): + torch.manual_seed(0) + print(f"device: {torch.cuda.get_device_name(0)}\n", flush=True) + + shapes = [ + ("matmul N=32 D=64 M=16 ", (32, 64, 16)), + ("matmul N=32 D=128 M=64 ", (32, 128, 64)), + ("matmul N=128 D=1152 M=1152", (128, 1152, 1152)), + ("GEMV N=1 D=64 M=16 ", (1, 64, 16)), + ("GEMV N=1 D=128 M=64 ", (1, 128, 64)), + ("GEMV N=1 D=1152 M=1152", (1, 1152, 1152)), + ] + + print("ORIGINAL scmp_llm (q_clip = q_norm - 2 = 125 for 8-bit)") + print("-" * 110, flush=True) + inputs = [] + for name, (N, D, M) in shapes: + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + inputs.append((name, a, b)) + run_pair(name, a, b, patch=False) + + print() + print("PATCHED scmp_llm (q_clip = q_norm = 127 ⇒ should match scmp_kernels)") + print("-" * 110, flush=True) + for name, a, b in inputs: + run_pair(name, a, b, patch=True) + + +if __name__ == "__main__": + main() diff --git a/tools/compare_old_vs_new_thorough.py b/tools/compare_old_vs_new_thorough.py new file mode 100644 index 0000000..faa3e5b --- /dev/null +++ b/tools/compare_old_vs_new_thorough.py @@ -0,0 +1,167 @@ +"""Exhaustive re-verification: scmp_llm/SC vs scmp_kernels.sc.sc_matmul. + +For each (granularity, mode) combination, sweeps: + - 3 different shapes + - 3 different random seeds + - 3 different stoc_len values (64, 128, 256) + - 2 different sc_prec values (6 and 8) where applicable + +Every case must report BIT-IDENTICAL (max|Δ|=0). Anything else means a real +divergence remains. +""" +from __future__ import annotations +import sys +from pathlib import Path + +# Old impl +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_triton as old_sc +from config_helpers import make_sobol_simple_config + +# New impl +from scmp_kernels.sc import sc_matmul as new_sc_matmul, clear_rng_cache as new_clear + +import torch + + +# Sanity: prove the two impls are genuinely separate modules +print(f"old module: {old_sc.__file__}", flush=True) +print(f"old enable fn: id={id(old_sc.sc_matmul_enable_triton)}", flush=True) + +import scmp_kernels.sc.kernels as new_k +print(f"new module: {new_k.__file__}", flush=True) +print(f"new dispatch: id={id(new_sc_matmul)}", flush=True) + +assert old_sc.__file__ != new_k.__file__, \ + "scmp_llm and scmp_kernels share a file — comparison would be vacuous" +print() + + +def reset(): + old_sc.clear_rng_cache() + new_clear() + + +fail_count = 0 +case_count = 0 + + +def case(label, a, b, *, granularity, mode, sc_prec, stoc_len, + group_a=1, group_b=1, old_kind="enable_triton"): + """Run both, compare, increment counters.""" + global fail_count, case_count + case_count += 1 + D = a.shape[-1] + cfg = make_sobol_simple_config(D, D, sc_prec) + + reset() + if old_kind == "enable_triton": + out_old = old_sc.sc_matmul_enable_triton( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + elif old_kind == "enable_batched_bipolar": + out_old = old_sc.sc_matmul_enable_batched_bipolar( + a, b, a.amax(dim=(1, 2)), a.amin(dim=(1, 2)), + b.amax(dim=(1, 2)), b.amin(dim=(1, 2)), + sc_prec, cfg, stoc_len=stoc_len, + ) + elif old_kind == "grouped_enable": + out_old = old_sc.sc_matmul_grouped_enable_triton( + a, b, group_a=group_a, group_b=group_b, + mode=mode, sc_prec=sc_prec, config=cfg, stoc_len=stoc_len, + ) + else: + raise ValueError(old_kind) + + reset() + kw = dict(granularity=granularity, mode=mode, + sc_prec=sc_prec, stoc_len=stoc_len, config=cfg) + if granularity == "per_row": + kw.update(group_a=group_a, group_b=group_b) + out_new = new_sc_matmul(a, b, **kw) + + max_abs = float((out_old - out_new).abs().max().item()) + ok = (max_abs == 0.0) + if not ok: + fail_count += 1 + flag = "✓" if ok else "✗" + print(f" {flag} {label:<70} max|Δ|={max_abs:.3e}", flush=True) + + +def main(): + print(f"device: {torch.cuda.get_device_name(0)}", flush=True) + print(f"\n{'─'*100}\nBIPOLAR 2D per_tensor ─ the path where the clipping fix matters most", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (N, D, M) in [(8, 32, 8), (32, 128, 64), (128, 1152, 1152)]: + for sc_prec, stoc_len in [(8, 64), (8, 128), (8, 256), (6, 64)]: + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + label = f"seed={seed} N={N:<4} D={D:<5} M={M:<5} sc_prec={sc_prec} stoc_len={stoc_len:<3}" + case(label, a, b, granularity="per_tensor", mode="bipolar", + sc_prec=sc_prec, stoc_len=stoc_len) + + print(f"\n{'─'*100}\nBIPOLAR GEMV per_tensor (N=1)", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (D, M) in [(32, 8), (128, 64), (1152, 1152)]: + for stoc_len in [64, 128, 256]: + a = torch.randn(1, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + label = f"seed={seed} GEMV D={D:<5} M={M:<5} sc_prec=8 stoc_len={stoc_len:<3}" + case(label, a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=stoc_len) + + print(f"\n{'─'*100}\nUNIPOLAR 2D per_tensor", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (N, D, M) in [(8, 32, 8), (32, 128, 64)]: + for stoc_len in [128, 256]: + a = torch.rand(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + label = f"seed={seed} unipolar N={N:<4} D={D:<5} M={M:<5} stoc_len={stoc_len:<3}" + case(label, a, b, granularity="per_tensor", mode="unipolar", + sc_prec=8, stoc_len=stoc_len) + + print(f"\n{'─'*100}\nQK per_head bipolar (3D batched)", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (BH, N, D) in [(4, 16, 32), (8, 64, 64), (16, 256, 72)]: + for stoc_len in [128, 256]: + q = torch.randn(BH, N, D, device="cuda") + k = torch.randn(BH, N, D, device="cuda") + label = f"seed={seed} QK BH={BH:<3} N={N:<5} D={D:<5} stoc_len={stoc_len:<3}" + case(label, q, k, granularity="per_head", mode="bipolar", + sc_prec=8, stoc_len=stoc_len, old_kind="enable_batched_bipolar") + + print(f"\n{'─'*100}\nAV grouped per_row bipolar", flush=True) + print("─" * 100, flush=True) + for seed in [0, 1, 42]: + torch.manual_seed(seed) + for (N, D, M) in [(16, 8, 32), (64, 32, 128), (256, 72, 256)]: + for stoc_len in [128, 256]: + attn = torch.softmax(torch.randn(N, M, device="cuda"), dim=-1) + v = torch.randn(D, M, device="cuda") * 0.1 + label = f"seed={seed} AV N={N:<4} D={D:<4} M={M:<4} stoc_len={stoc_len:<3}" + case(label, attn, v, granularity="per_row", mode="bipolar", + sc_prec=8, stoc_len=stoc_len, + group_a=N, group_b=D, old_kind="grouped_enable") + + print(f"\n{'='*100}", flush=True) + print(f"Total cases: {case_count} Bit-identical: {case_count - fail_count} Diverged: {fail_count}", flush=True) + print("=" * 100, flush=True) + if fail_count == 0: + print("\n ✓ ALL CASES BIT-IDENTICAL — scmp_llm and scmp_kernels produce equivalent SC outputs.", flush=True) + else: + print(f"\n ✗ {fail_count} cases diverged — investigate.", flush=True) + sys.exit(0 if fail_count == 0 else 1) + + +if __name__ == "__main__": + main() diff --git a/tools/compare_vs_cpu_reference.py b/tools/compare_vs_cpu_reference.py new file mode 100644 index 0000000..0330a96 --- /dev/null +++ b/tools/compare_vs_cpu_reference.py @@ -0,0 +1,156 @@ +"""Compare scmp_kernels.sc.sc_matmul (GPU Triton) against the CPU reference +in scmp_llm/SC/sc_enable.py (sc_matmul_enable). + +The CPU reference is the gold-standard semantic check: it uses Python/PyTorch +ops, no Triton, two algorithms: + - cycle_by_cycle: exact UnarySim FSUMul simulator (very slow, small shapes) + - k_shortcut: vectorized prefix-sum equivalent (faster) + +For each shape × mode × stoc_len, runs: + GPU Triton (new scmp_kernels) + CPU k_shortcut (scmp_llm reference) + CPU cycle_by_cycle (scmp_llm reference, on small shapes only) + +Compares each pair. Expectations: + - CPU k_shortcut vs cycle_by_cycle: should be bit-identical (same math). + - GPU vs CPU: should be bit-identical OR differ only by float32 rounding + (≤ 1e-4 absolute, dominated by accumulation order). +""" +from __future__ import annotations +import sys, time +from pathlib import Path + +SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") +sys.path.insert(0, str(SCMP_LLM_SC)) +import sc_enable as cpu_ref +from config_helpers import make_sobol_simple_config + +from scmp_kernels.sc import sc_matmul as gpu_matmul, clear_rng_cache + +import torch + + +def diff(a, b): + d = (a - b).abs() + rel = d / b.abs().clamp_min(1e-6) + return float(d.max().item()), float(d.mean().item()), float(rel.max().item()) + + +def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + +def reset(): + clear_rng_cache() + + +def section(title): + print(f"\n{title}\n{'─'*120}", flush=True) + print(f" {'CASE':<55} {'max|Δ|':>10} {'mean|Δ|':>10} {'max rel':>10} {'verdict':>20}", flush=True) + + +def verdict(ma, threshold=1e-4): + if ma == 0.0: return "BIT-IDENTICAL" + if ma < threshold: return f"≈identical ({ma:.1e})" + return f"differ ({ma:.2e})" + + +def main(): + print(f"device: {torch.cuda.get_device_name(0)}", flush=True) + print(f"CPU reference: {cpu_ref.__file__}", flush=True) + print(f"GPU impl: scmp_kernels.sc.sc_matmul", flush=True) + + # ---------------------------------------------------------------- + # 1. CPU k_shortcut vs CPU cycle_by_cycle (both reference) + # ---------------------------------------------------------------- + section("1. CPU references self-consistent? — k_shortcut vs cycle_by_cycle") + torch.manual_seed(0) + for (N, D, M) in [(4, 16, 4), (8, 32, 8)]: # cycle_by_cycle is slow + for mode in ["bipolar", "unipolar"]: + if mode == "unipolar": + a = torch.rand(N, D, device="cuda") + else: + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + cfg = make_sobol_simple_config(D, D, 8) + out_cyc = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=8, config=cfg, method="cycle_by_cycle") + out_ks = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode=mode, sc_prec=8, config=cfg, method="k_shortcut") + ma, me, mr = diff(out_ks, out_cyc) + print(f" {f'{mode:<8} N={N:<3} D={D:<4} M={M:<3} sc_prec=8 stoc_len=256':<55} " + f"{ma:>10.3e} {me:>10.3e} {mr:>10.3e} {verdict(ma):>20}", flush=True) + + # ---------------------------------------------------------------- + # 2. GPU Triton vs CPU k_shortcut — bipolar per_tensor + # ---------------------------------------------------------------- + section("2. GPU Triton vs CPU k_shortcut — bipolar per_tensor") + for (N, D, M) in [(8, 32, 8), (32, 128, 64), (64, 256, 32)]: + for stoc_len in [64, 256]: + torch.manual_seed(0) + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + cfg = make_sobol_simple_config(D, D, 8) + reset() + out_cpu = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="bipolar", sc_prec=8, config=cfg, method="k_shortcut") + # CPU ref uses stoc_len = 2**sc_prec = 256 always (no stoc_len arg!) + # So we only compare at stoc_len=256 for full equivalence + if stoc_len != 256: + continue + reset() + out_gpu = gpu_matmul(a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=stoc_len, config=cfg) + ma, me, mr = diff(out_gpu, out_cpu) + print(f" {f'bipolar N={N:<3} D={D:<4} M={M:<3} stoc_len={stoc_len:<3}':<55} " + f"{ma:>10.3e} {me:>10.3e} {mr:>10.3e} {verdict(ma):>20}", flush=True) + + # ---------------------------------------------------------------- + # 3. GPU vs CPU — unipolar + # ---------------------------------------------------------------- + section("3. GPU Triton vs CPU k_shortcut — unipolar per_tensor") + for (N, D, M) in [(8, 32, 8), (16, 128, 32)]: + torch.manual_seed(0) + a = torch.rand(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + cfg = make_sobol_simple_config(D, D, 8) + reset() + out_cpu = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="unipolar", sc_prec=8, config=cfg, method="k_shortcut") + reset() + out_gpu = gpu_matmul(a, b, granularity="per_tensor", mode="unipolar", + sc_prec=8, stoc_len=256, config=cfg) + ma, me, mr = diff(out_gpu, out_cpu) + print(f" {f'unipolar N={N:<3} D={D:<4} M={M:<3} stoc_len=256':<55} " + f"{ma:>10.3e} {me:>10.3e} {mr:>10.3e} {verdict(ma):>20}", flush=True) + + # ---------------------------------------------------------------- + # 4. Numerical accuracy: both vs torch.matmul fp baseline + # ---------------------------------------------------------------- + section("4. SC vs fp baseline — both impls should match torch.matmul within SC noise band") + torch.manual_seed(0) + N, D, M = 32, 128, 64 + a = torch.randn(N, D, device="cuda") + b = torch.randn(M, D, device="cuda") * 0.1 + fp = a @ b.t() + cfg = make_sobol_simple_config(D, D, 8) + reset() + out_gpu = gpu_matmul(a, b, granularity="per_tensor", mode="bipolar", + sc_prec=8, stoc_len=256, config=cfg) + out_cpu = cpu_ref.sc_matmul_enable( + a, b, a.max().item(), a.min().item(), b.max().item(), b.min().item(), + mode="bipolar", sc_prec=8, config=cfg, method="k_shortcut") + print(f" bipolar N=32 D=128 M=64 stoc_len=256", flush=True) + print(f" rel_err GPU vs fp: {rel_err(out_gpu, fp):.4f}", flush=True) + print(f" rel_err CPU vs fp: {rel_err(out_cpu, fp):.4f}", flush=True) + print(f" rel_err GPU vs CPU: {rel_err(out_gpu, out_cpu):.4e}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/sc_call_counter.py b/tools/sc_call_counter.py new file mode 100644 index 0000000..ab59822 --- /dev/null +++ b/tools/sc_call_counter.py @@ -0,0 +1,61 @@ +"""Wrap sc_matmul to count calls per granularity, then run quant_sc_main.py. + +Confirms SC kernels are actually invoked during DiT sampling, not skipped. +Prints a count summary at process exit. +""" +import sys, os, atexit +from pathlib import Path +from collections import Counter + +ROOT = Path('/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion') +sys.path.insert(0, str(ROOT)) + +import scmp_kernels.sc.matmul as mm +_orig = mm.sc_matmul +counts = Counter() + +def _wrapped(a, b, granularity="per_row", **kw): + key = ( + granularity, + kw.get('mode', 'bipolar'), + 'chunk' if kw.get('chunk_d', 0) > 0 else 'nochunk', + 'group' if (kw.get('group_a', 1) > 1 or kw.get('group_b', 1) > 1) else 'nogroup', + ) + counts[key] += 1 + return _orig(a, b, granularity=granularity, **kw) + +mm.sc_matmul = _wrapped + +# Patch every site that did `from scmp_kernels.sc import sc_matmul` +import scmp_kernels.sc as sc_pkg +sc_pkg.sc_matmul = _wrapped +import qdit.sc_integration.sc_attention as sa +import qdit.sc_integration.sc_mlp as sm +sa.sc_matmul = _wrapped +sm.sc_matmul = _wrapped + +def _summary(): + total = sum(counts.values()) + print(f"\n========== sc_matmul call summary ==========", flush=True) + print(f"Total sc_matmul invocations: {total}", flush=True) + for k, n in sorted(counts.items(), key=lambda kv: -kv[1]): + gran, mode, chunk, group = k + print(f" {n:>6} granularity={gran:<10} mode={mode:<8} {chunk:<8} {group}", flush=True) + print(f"============================================\n", flush=True) +atexit.register(_summary) + +# Now run quant_sc_main as if invoked directly +sys.argv = [ + 'quant_sc_main.py', + '--ckpt', '/nfs/turbo/coe-nbleier/zhkangqi/pretrained_models/DiT-XL-2-256x256.pt', + '--wbits', '8', '--abits', '8', '--w_sym', '--a_sym', + '--timewise', '0.5', '--qklayerwise', '0.8', + '--avlayerwise', '0.0', '--projlayerwise', '0.0', + '--mlplayerwise', '0.0', '--inputprojlayerwise', '0.0', + '--sc_prec', '8', + '--image-size', '256', '--num-sampling-steps', '50', + '--cfg-scale', '4', '--batch-size', '8', + '--results-dir', 'results/smoke_e2e_count', +] +os.chdir(str(ROOT)) +exec(compile(open('scripts/quant_sc_main.py').read(), 'scripts/quant_sc_main.py', 'exec')) diff --git a/tools/sc_call_counter_full.py b/tools/sc_call_counter_full.py new file mode 100644 index 0000000..197738c --- /dev/null +++ b/tools/sc_call_counter_full.py @@ -0,0 +1,59 @@ +"""Full SC: all operators enabled, uniform stoc_len=128. Count sc_matmul calls.""" +import sys, os, atexit +from pathlib import Path +from collections import Counter + +ROOT = Path('/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion') +sys.path.insert(0, str(ROOT)) + +import scmp_kernels.sc.matmul as mm +_orig = mm.sc_matmul +counts = Counter() + +def _wrapped(a, b, granularity="per_row", **kw): + key = ( + granularity, + kw.get('mode', 'bipolar'), + 'chunk' if kw.get('chunk_d', 0) > 0 else 'nochunk', + 'group' if (kw.get('group_a', 1) > 1 or kw.get('group_b', 1) > 1) else 'nogroup', + ) + counts[key] += 1 + return _orig(a, b, granularity=granularity, **kw) + +mm.sc_matmul = _wrapped +import scmp_kernels.sc as sc_pkg +sc_pkg.sc_matmul = _wrapped +import qdit.sc_integration.sc_attention as sa +import qdit.sc_integration.sc_mlp as sm +sa.sc_matmul = _wrapped +sm.sc_matmul = _wrapped + +def _summary(): + total = sum(counts.values()) + print(f"\n========== sc_matmul call summary (uniform stoc_len=128, ALL OPS) ==========", flush=True) + print(f"Total sc_matmul invocations: {total}", flush=True) + for k, n in sorted(counts.items(), key=lambda kv: -kv[1]): + gran, mode, chunk, group = k + print(f" {n:>6} granularity={gran:<10} mode={mode:<8} {chunk:<8} {group}", flush=True) + print(f"============================================================================\n", flush=True) +atexit.register(_summary) + +sys.argv = [ + 'quant_sc_main.py', + '--ckpt', '/nfs/turbo/coe-nbleier/zhkangqi/pretrained_models/DiT-XL-2-256x256.pt', + '--wbits', '8', '--abits', '8', '--w_sym', '--a_sym', + '--timewise', '1.0', + '--qklayerwise', '1.0', + '--avlayerwise', '1.0', + '--projlayerwise', '1.0', + '--mlplayerwise', '1.0', + '--inputprojlayerwise', '1.0', + '--sc_prec', '8', + '--sc_fixed_level_prec', + '--sc_config', 'results/sc_cfg_uniform128_all.json', + '--image-size', '256', '--num-sampling-steps', '50', + '--cfg-scale', '4', '--batch-size', '8', + '--results-dir', 'results/smoke_e2e_all_L128', +] +os.chdir(str(ROOT)) +exec(compile(open('scripts/quant_sc_main.py').read(), 'scripts/quant_sc_main.py', 'exec')) diff --git a/tools/smoke_test_e2e.py b/tools/smoke_test_e2e.py new file mode 100644 index 0000000..3ee3700 --- /dev/null +++ b/tools/smoke_test_e2e.py @@ -0,0 +1,223 @@ +"""End-to-end smoke test for the scmp_diffusion bootstrap. + +Run on a GPU node. Verifies: + 1. submodule imports resolve (scmp_kernels.sc / scmp_kernels.mp) + 2. sc_matmul reaches every granularity + mode combination + 3. SC results stay within sane rel-err of torch.matmul + 4. Q-DiT integration imports resolve (qdit.sc_integration.{sc_attention, sc_mlp, noise_matmul}) + 5. SCMlp forward pass produces a finite, non-zero, sensibly-scaled output + 6. SCAttention forward pass produces a finite output + 7. No remaining references to deprecated names + +Exits non-zero on any failure with a clear message. +""" +from __future__ import annotations + +import importlib +import sys +import time +import traceback +from pathlib import Path + +# Make local Q-DiT package importable +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + + +def section(name): + print(f"\n{'='*80}\n{name}\n{'='*80}", flush=True) + + +def check(label, fn): + t0 = time.time() + try: + result = fn() + dt = (time.time() - t0) * 1000 + msg = f" PASS {label:<60} ({dt:>6.1f} ms)" + if result: + msg += f" {result}" + print(msg, flush=True) + return True + except Exception: + print(f" FAIL {label}", flush=True) + traceback.print_exc() + return False + + +def main(): + failures = [] + + section("1. Imports") + import torch + print(f" torch {torch.__version__}, cuda available: {torch.cuda.is_available()}", flush=True) + if not torch.cuda.is_available(): + print(" no CUDA — skipping kernel execution"); sys.exit(2) + print(f" device: {torch.cuda.get_device_name(0)}", flush=True) + try: + import triton + print(f" triton {triton.__version__}", flush=True) + except Exception as e: + print(f" triton import: {e}", flush=True) + sys.exit(2) + + if not check("import scmp_kernels.sc", + lambda: importlib.import_module("scmp_kernels.sc").__name__): + failures.append("scmp_kernels.sc") + + if not check("import scmp_kernels.mp", + lambda: importlib.import_module("scmp_kernels.mp").__name__): + failures.append("scmp_kernels.mp") + + if not check("from scmp_kernels.sc import sc_matmul, clear_rng_cache, det_kernel_tuning", + lambda: __import__("scmp_kernels.sc", fromlist=["sc_matmul","clear_rng_cache","det_kernel_tuning"])): + failures.append("sc public API") + + if not check("from scmp_kernels.mp import (9 names)", + lambda: [getattr(__import__("scmp_kernels.mp", fromlist=[n]), n) + for n in ("MPConfig","AdaptiveMPConfig","RangeMPConfig","RowAssignment", + "classify_rows_by_metric","adaptive_classify_rows", + "classify_groups_by_range","MPDistributionLogger","MetricProfiler")] + and "OK"): + failures.append("mp public API") + + section("2. sc_matmul granularity sweep — basic correctness vs torch.matmul") + from scmp_kernels.sc import sc_matmul + from scmp_kernels.sc.config_helpers import make_sobol_simple_config + + def rel_err(pred, target): + num = (pred - target).pow(2).mean().sqrt() + den = target.pow(2).mean().sqrt().clamp_min(1e-8) + return float((num / den).item()) + + torch.manual_seed(0) + device = "cuda" + + # per_tensor 2D bipolar + def per_tensor_2d_bipolar(): + a = torch.randn(32, 128, device=device, dtype=torch.float32) + b = torch.randn(64, 128, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_tensor", mode="bipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_tensor 2D bipolar (32×128)@(64×128).T", per_tensor_2d_bipolar): + failures.append("per_tensor 2D bipolar") + + def per_row_2d_bipolar(): + a = torch.randn(32, 128, device=device, dtype=torch.float32) + b = torch.randn(64, 128, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="bipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_row 2D bipolar", per_row_2d_bipolar): + failures.append("per_row 2D bipolar") + + def per_row_mlp_chunked(): + a = torch.randn(32, 1152, device=device, dtype=torch.float32) + b = torch.randn(64, 1152, device=device, dtype=torch.float32) * 0.05 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="bipolar", + chunk_d=72, sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)} (chunked D=1152→72)" + + if not check("sc_matmul per_row MLP chunked", per_row_mlp_chunked): + failures.append("per_row MLP chunked") + + def per_row_grouped(): + N, M, D = 64, 32, 128 + a = torch.softmax(torch.randn(N, M, device=device), dim=-1) + b = torch.randn(D, M, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="bipolar", + group_a=N, group_b=D, sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)} (group_a={N},group_b={D})" + + if not check("sc_matmul per_row grouped (AV pattern)", per_row_grouped): + failures.append("per_row grouped") + + def per_head_bipolar(): + BH, N, D = 8, 32, 64 + q = torch.randn(BH, N, D, device=device, dtype=torch.float32) + k = torch.randn(BH, N, D, device=device, dtype=torch.float32) + fp = q @ k.transpose(-1, -2) + sc = sc_matmul(q, k, granularity="per_head", mode="bipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_head bipolar (8×32×64)@(8×32×64).T", per_head_bipolar): + failures.append("per_head bipolar") + + def per_row_unipolar(): + a = torch.rand(32, 64, device=device, dtype=torch.float32) + b = torch.randn(16, 64, device=device, dtype=torch.float32) * 0.1 + fp = a @ b.t() + sc = sc_matmul(a, b, granularity="per_row", mode="unipolar", sc_prec=8, stoc_len=256) + e = rel_err(sc, fp) + return f"rel_err={e:.4f} shape={tuple(sc.shape)}" + + if not check("sc_matmul per_row 2D unipolar", per_row_unipolar): + failures.append("per_row unipolar") + + section("3. Q-DiT integration imports") + + qdit_imports = [ + ("qdit.sc_integration", ["SCController", "SCAttention", "SCMlp", "SCDiTBlock", "MPConfig", "add_sc_wrapper", "create_sc_controller_from_args"]), + ("qdit.sc_integration.sc_attention", ["SCAttention"]), + ("qdit.sc_integration.sc_mlp", ["SCMlp"]), + ("qdit.sc_integration.noise_matmul", ["noisy_sc_matmul"]), + ("qdit.sc_integration.sc_controller", ["SCController"]), + ("qdit.sc_integration.mp_config", ["MPConfig", "AdaptiveMPConfig"]), + ] + + for mod_name, names in qdit_imports: + def _doit(m=mod_name, ns=names): + mod = importlib.import_module(m) + missing = [n for n in ns if not hasattr(mod, n)] + if missing: + raise AttributeError(f"missing names in {m}: {missing}") + return f"resolved: {', '.join(ns)}" + if not check(f"import {mod_name}", _doit): + failures.append(mod_name) + + section("4. Deprecated-name surveillance") + + deprecated = ["sc_matmul_per_tensor", "sc_matmul_mlp", "sc_matmul_grouped", + "sc_matmul_enable_triton", "sc_matmul_enable_triton_mlp", + "sc_matmul_grouped_enable_triton", "sc_matmul_enable_batched_bipolar", + "bin_to_stoc_packed", "xnor_matmul"] + def deprecated_absent(): + import scmp_kernels.sc as sck_sc + present = [n for n in deprecated if hasattr(sck_sc, n)] + if present: + raise AssertionError(f"deprecated names still exported: {present}") + return f"none of {len(deprecated)} deprecated names re-emerged" + if not check("scmp_kernels.sc has no deprecated public names", deprecated_absent): + failures.append("deprecated surveillance") + + def sc_enable_absent(): + import qdit.sc_integration.sc_controller as ctrl + # constructor should not accept sc_enable + from inspect import signature + params = signature(ctrl.SCController.__init__).parameters + if "sc_enable" in params: + raise AssertionError("SCController.__init__ still accepts sc_enable") + return "SCController has no sc_enable parameter" + if not check("SCController no longer accepts sc_enable", sc_enable_absent): + failures.append("sc_enable removed") + + section("Summary") + if failures: + print(f"\n {len(failures)} FAILURES:") + for f in failures: print(f" - {f}") + sys.exit(1) + print("\n ALL CHECKS PASSED") + sys.exit(0) + + +if __name__ == "__main__": + main() From ca5eab07f2d74f9a6f998b0e5571f151b0efa77c Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:55:47 -0400 Subject: [PATCH 4/6] kernels: bump submodule to trace(#25)+#28 scramble-cache fix + doc cleanup Submodule now at local branch local/trace+scramble-cache (6c4539a): PR#25 trace subsystem + cherry-picked PR#28 enable-table scramble cache-key fix + mp/config doc cleanup. NOTE: 6c4539a is LOCAL-ONLY (not yet pushed to the heroarmor fork); a fresh 'git submodule update' cannot resolve it until pushed. Co-Authored-By: Claude Opus 4.8 (1M context) --- scmp_kernels | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scmp_kernels b/scmp_kernels index 9358239..6c4539a 160000 --- a/scmp_kernels +++ b/scmp_kernels @@ -1 +1 @@ -Subproject commit 9358239af8d65f54a08e4a47e090933947626a41 +Subproject commit 6c4539ae35a60ce95f3e7747edc81444d823cc55 From 1b74fab82b16ba3ebf0f3b887744d6c823cbead2 Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Wed, 15 Jul 2026 04:58:22 -0400 Subject: [PATCH 5/6] fix(halve): enforce max stoc_len 128 end-to-end; cfg-matched calibration; mse_rel metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The halve MP line calibrated and ran levels above the halve-mode maximum (2^(sc_prec-1)=128): the 256/192 levels are unrealizable on halve hardware and silently inflated attention accuracy (up to 99% of qk rows ran 256/192- cycle streams). Fixes, one per layer: - scmp_kernels bump -> c967d3a: sc_matmul now raises on explicit stoc_len/rng_levels above 2^(sc_prec-1) under halve+bipolar (upstream kernels PR #30), so stale 256-max tables crash instead of overspending. - calibrate_mp_thresholds.py: hard-validate mp_levels and budget_ref_stoc_len <= 2^(sc_prec-1) under --sc_halve; add --metric mse_rel (squared relative L2 = cosine's quadratic curvature + gain-error sensitivity; a pure 10% gain error scores 0.0 under cosine, 0.01 under mse_rel). - sbatch_calib_globalpr_flop_halve.sb / sbatch_calib_globalpr_halve.sb: 128-capped level grid, budget_ref 128, BUDGET_RATIO computed from AVG in-script (avg/128 — passing the old avg/256 ratios silently halved every budget), optional METRIC arg with suffixed OUTDIR, and calibration now runs at the deployment CFG (teacher_cfg_scale 1.5, batch 16 so the CFG-doubled batch keeps the old 32-row memory footprint) instead of cfg=0. - sbatch_mp_globalpr_flop_halve_auto.sb: [METRIC] and [NUM_FID] args (metric-suffixed calib/output dirs; NUM_FID for fast probes), points at the 128-capped tables; old 256-grid tables/arms preserved in the un-suffixed dirs for comparison. Validated: avg48 recalibration hits expected_avg_stoc_len=48.00; runtime realized per-op distribution matches the table within 1-3pp per level; on the noise-aligned index prefix the fixed adaptive arms (cosine and mse_rel) beat uniform48 on fidelity-to-FP (2.5% vs 5% trajectory divergence). Co-Authored-By: Claude Fable 5 --- scmp_kernels | 2 +- scripts/calibrate_mp_thresholds.py | 41 ++++++++++++++++++- scripts/sbatch_calib_globalpr_flop_halve.sb | 30 ++++++++++---- scripts/sbatch_calib_globalpr_halve.sb | 26 ++++++++---- scripts/sbatch_mp_globalpr_flop_halve_auto.sb | 22 ++++++---- 5 files changed, 95 insertions(+), 26 deletions(-) diff --git a/scmp_kernels b/scmp_kernels index 6c4539a..c967d3a 160000 --- a/scmp_kernels +++ b/scmp_kernels @@ -1 +1 @@ -Subproject commit 6c4539ae35a60ce95f3e7747edc81444d823cc55 +Subproject commit c967d3aa9014f37a2f6823048dd02993e2e1dd21 diff --git a/scripts/calibrate_mp_thresholds.py b/scripts/calibrate_mp_thresholds.py index 4b70425..92f5596 100644 --- a/scripts/calibrate_mp_thresholds.py +++ b/scripts/calibrate_mp_thresholds.py @@ -141,6 +141,22 @@ def _cosine_dist_heads(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor return 1.0 - cos +def _relative_mse_rows(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Squared relative L2 per row: ||d||^2 / ||t||^2. + + Small-error expansion = gain_err^2 + orthogonal_r^2: keeps cosine's + quadratic (second-order-loss) curvature and relative row weighting, but + also sees pure gain/bias errors that cosine is blind to. + """ + r = _relative_l2_rows(pred, target) + return r * r + + +def _relative_mse_heads(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + r = _relative_l2_heads(pred, target) + return r * r + + # Globals dispatched from CLI flags --metric / --teacher. _METRIC_ROWS = _relative_l2_rows _METRIC_HEADS = _relative_l2_heads @@ -1098,9 +1114,10 @@ def _build_parser(): parser.add_argument( "--metric", type=str, - choices=["l2", "cosine"], + choices=["l2", "cosine", "mse_rel"], default="cosine", - help="Per-unit error metric. cosine = 1 - cos_sim (recommended for fix mode).", + help="Per-unit error metric. cosine = 1 - cos_sim (recommended for fix mode); " + "mse_rel = ||d||^2/||t||^2 (cosine's curvature + gain-error sensitivity).", ) parser.add_argument( "--teacher", @@ -1127,6 +1144,23 @@ def main(): explicit_timesteps = _parse_csv_ints(args.calib_timesteps) budget_ref_stoc_len = args.budget_ref_stoc_len or max(levels) + # Halve mode: the sign-magnitude grid (and hence the longest realizable + # stream) is 2^(sc_prec-1). Levels above that would probe accuracy the + # halve hardware cannot deliver, and a larger budget_ref silently rescales + # every budget_ratio, so reject both loudly. + if args.sc_halve: + halve_max = 2 ** (args.sc_prec - 1) + bad = [lv for lv in levels if lv > halve_max] + if bad: + raise ValueError( + f"--sc_halve: mp_levels {bad} exceed the halve-mode maximum " + f"2^(sc_prec-1)={halve_max}; use levels <= {halve_max}.") + if budget_ref_stoc_len > halve_max: + raise ValueError( + f"--sc_halve: budget_ref_stoc_len={budget_ref_stoc_len} exceeds " + f"the halve-mode maximum {halve_max}. Express budget_ratio " + f"against {halve_max} (avg = budget_ratio * {halve_max}).") + latent_size = args.image_size // 8 model = DiT_models[args.model]( input_size=latent_size, @@ -1168,6 +1202,9 @@ def main(): if args.metric == "cosine": _METRIC_ROWS = _cosine_dist_rows _METRIC_HEADS = _cosine_dist_heads + elif args.metric == "mse_rel": + _METRIC_ROWS = _relative_mse_rows + _METRIC_HEADS = _relative_mse_heads _USE_FP_TEACHER = (args.teacher == "fp") _SC_HALVE = bool(args.sc_halve) print(f"SC halve during calibration: {_SC_HALVE}") diff --git a/scripts/sbatch_calib_globalpr_flop_halve.sb b/scripts/sbatch_calib_globalpr_flop_halve.sb index 06d1984..4239749 100644 --- a/scripts/sbatch_calib_globalpr_flop_halve.sb +++ b/scripts/sbatch_calib_globalpr_flop_halve.sb @@ -8,26 +8,40 @@ #SBATCH --time=03:00:00 # FLOP-weighted + HALVE per-row GLOBAL calibration for one budget (fix branch: # fix/flop-cost-perrow-qk-drop-alphabeta). Matches the halve MP deployment. -# Usage: sbatch sbatch_calib_globalpr_flop_halve.sb +# Halve mode caps stream length at 2^(sc_prec-1)=128: levels above 128 are +# unrealizable on halve hardware, so the level grid tops out at 128 and the +# budget is expressed against a 128 reference (BUDGET_RATIO = AVG/128, +# computed here — do NOT pass the old AVG/256 ratios). +# Usage: sbatch sbatch_calib_globalpr_flop_halve.sb [METRIC] +# METRIC: cosine (default) | mse_rel | l2 — non-cosine tables land in a +# metric-suffixed OUTDIR so A/B arms never overwrite each other. set -uo pipefail -AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +AVG="${1:?AVG}" +METRIC="${2:-cosine}" +HALVE_REF=128 +if (( AVG > HALVE_REF )); then + echo "ERROR: AVG=${AVG} exceeds halve-mode max stoc_len ${HALVE_REF}" >&2 + exit 1 +fi +BR=$(awk "BEGIN{print ${AVG}/${HALVE_REF}}") REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt -OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve_max128 +[[ "$METRIC" != "cosine" ]] && OUTDIR="${OUTDIR}_${METRIC}" source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" -echo "=== FLOP+HALVE calib GLOBAL per-row avg${AVG} br=${BR} branch=$(git branch --show-current) $(date) ===" +echo "=== FLOP+HALVE calib GLOBAL per-row avg${AVG} br=${BR} (ref=${HALVE_REF}) metric=${METRIC} branch=$(git branch --show-current) $(date) ===" python -u scripts/calibrate_mp_thresholds.py \ - --mp_levels 256,192,128,96,64,48,32,16 \ - --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ - --metric cosine --teacher fp \ + --mp_levels 128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len "$HALVE_REF" --budget_scope global \ + --metric "$METRIC" --teacher fp \ --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ --wbits 8 --abits 8 --w_sym --a_sym \ --image-size 256 --num-sampling-steps 50 \ --num_calib_batches 1 --num_calib_timesteps 50 \ --timestep_buckets 10 --layer_buckets 4 \ - --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ + --sc_halve --teacher_cfg_scale 1.5 --batch-size 16 --ckpt "$CKPT" \ --calib_output_json "$OUTDIR/calib_fix_avg${AVG}.json" \ --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_calib_globalpr_halve.sb b/scripts/sbatch_calib_globalpr_halve.sb index d2f0403..139eca2 100755 --- a/scripts/sbatch_calib_globalpr_halve.sb +++ b/scripts/sbatch_calib_globalpr_halve.sb @@ -6,26 +6,36 @@ #SBATCH --cpus-per-gpu=8 #SBATCH --mem=110G #SBATCH --time=03:00:00 -# per-row GLOBAL calibration for one budget. Usage: sbatch sbatch_calib_globalpr.sb +# per-row GLOBAL calibration (row-weighted) for one budget, HALVE mode. +# Halve mode caps stream length at 2^(sc_prec-1)=128, so the level grid tops +# out at 128 and the budget is expressed against a 128 reference +# (BUDGET_RATIO = AVG/128, computed here — do NOT pass the old AVG/256 ratios). +# Usage: sbatch sbatch_calib_globalpr_halve.sb set -uo pipefail -AVG="${1:?AVG}"; BR="${2:?BUDGET_RATIO}" +AVG="${1:?AVG}" +HALVE_REF=128 +if (( AVG > HALVE_REF )); then + echo "ERROR: AVG=${AVG} exceeds halve-mode max stoc_len ${HALVE_REF}" >&2 + exit 1 +fi +BR=$(awk "BEGIN{print ${AVG}/${HALVE_REF}}") REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt -OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_halve +OUTDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_halve_max128 source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit export PYTHONUNBUFFERED=1 SC_OWEN_MODE=bitrev; cd "$REPO"; mkdir -p "$OUTDIR" -echo "=== calib GLOBAL per-row HALVE avg${AVG} br=${BR} $(date) ===" +echo "=== calib GLOBAL per-row HALVE avg${AVG} br=${BR} (ref=${HALVE_REF}) $(date) ===" python -u scripts/calibrate_mp_thresholds.py \ - --mp_levels 256,192,128,96,64,48,32,16 \ - --budget_ratio "$BR" --budget_ref_stoc_len 256 --budget_scope global \ + --mp_levels 128,96,64,48,32,16 \ + --budget_ratio "$BR" --budget_ref_stoc_len "$HALVE_REF" --budget_scope global \ --metric cosine --teacher fp \ --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ --wbits 8 --abits 8 --w_sym --a_sym \ --image-size 256 --num-sampling-steps 50 \ --num_calib_batches 1 --num_calib_timesteps 50 \ --timestep_buckets 10 --layer_buckets 4 \ - --sc_halve --teacher_cfg_scale 0.0 --ckpt "$CKPT" \ - --calib_output_json "$OUTDIR/calib_fix_avg${AVG}_l256_ref192.json" \ + --sc_halve --teacher_cfg_scale 1.5 --batch-size 16 --ckpt "$CKPT" \ + --calib_output_json "$OUTDIR/calib_fix_avg${AVG}.json" \ --calib_summary_csv "$OUTDIR/calib_fix_avg${AVG}_summary.csv" 2>&1 | tee "$OUTDIR/calib_avg${AVG}.log" echo "DONE avg${AVG} $(date)" diff --git a/scripts/sbatch_mp_globalpr_flop_halve_auto.sb b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb index 7debbc4..d3f5b9a 100644 --- a/scripts/sbatch_mp_globalpr_flop_halve_auto.sb +++ b/scripts/sbatch_mp_globalpr_flop_halve_auto.sb @@ -9,18 +9,26 @@ # # FLOP-weighted + HALVE adaptive-MP generation, auto-resumes to NUM_FID, TIMEOUT-safe: # queues its successor (afterany, chain-capped) at the START. Uses the FLOP+halve -# calibrated tables. Usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] +# calibrated tables with the level grid capped at 128 (halve-mode max stoc_len = +# 2^(sc_prec-1); the old 256-max tables probed unrealizable levels — kept in +# calib_global_perrow_flop_halve/ + ..._flop_halve_cfg15/ for comparison only). +# Usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] [METRIC] [NUM_FID] +# METRIC: cosine (default) | mse_rel | l2 — must match the calib arm; picks the +# metric-suffixed calib table and output dir. +# NUM_FID: target image count (default 2000); e.g. 100 for a fast PSNR probe. set -uo pipefail -AVG="${1:?usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN]}"; CHAIN="${2:-0}"; NUM_GPUS=3; MAXCHAIN=15 +AVG="${1:?usage: sbatch sbatch_mp_globalpr_flop_halve_auto.sb [CHAIN] [METRIC] [NUM_FID]}"; CHAIN="${2:-0}"; METRIC="${3:-cosine}"; NUM_FID_ARG="${4:-2000}"; NUM_GPUS=3; MAXCHAIN=15 SELF=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion/scripts/sbatch_mp_globalpr_flop_halve_auto.sb REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt -CALIB=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg${AVG}.json -OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_flop_halve_cfg15/adaptive_avg${AVG} +CALIBDIR=$SCRATCH/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve_max128 +OUT=$SCRATCH/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/adaptive_avg${AVG} +if [[ "$METRIC" != "cosine" ]]; then CALIBDIR="${CALIBDIR}_${METRIC}"; OUT="${OUT}_${METRIC}"; fi +CALIB=$CALIBDIR/calib_fix_avg${AVG}.json SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs -NUM_FID=2000; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 -MP_LEVELS=256,192,128,96,64,48,32,16 +NUM_FID=$NUM_FID_ARG; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +MP_LEVELS=128,96,64,48,32,16 mkdir -p "$SAMPLES" "$IDX" "$LOG" [[ -f "$CALIB" ]] || { echo "ERROR: missing calib $CALIB" >&2; exit 1; } source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit @@ -32,7 +40,7 @@ if [[ "$START" -ge "$NUM_FID" ]]; then echo "[complete] avg${AVG} already at $NU # Queue successor NOW (afterany -> runs after me regardless of timeout/complete); chain-capped. if [[ "$CHAIN" -lt "$MAXCHAIN" ]]; then sbatch -o "$OUT/slurm-%j.out" -e "$OUT/slurm-%j.err" \ - --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) \ + --dependency=afterany:${SLURM_JOB_ID} "$SELF" "$AVG" $((CHAIN+1)) "$METRIC" "$NUM_FID_ARG" \ && echo "[chain] queued successor (chain=$((CHAIN+1)))" fi python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) From c328d36521448b54d0640701ff206b865bee110d Mon Sep 17 00:00:00 2001 From: heroarmor <162866837+heroarmor@users.noreply.github.com> Date: Wed, 15 Jul 2026 04:58:39 -0400 Subject: [PATCH 6/6] tools(probe): verification sbatch scripts for the halve MP line MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Small single-purpose probes developed while root-causing the halve-line regressions; kept because each verifies an invariant worth re-checking: - eval/sbatch_eval_uniform_halve.sb: KID + FID/IS/Prec/Recall for one uniform-halve arm (companion to sbatch_eval_mp_flop_halve.sb). - eval/sbatch_eval_ab_kid100.sb: matched-index KID across probe arms + an FP anchor on the same indices (the anchor exposes the class-composition floor that dominates absolute KID at small n — FP scores the same ~30x1e-3 as SC arms on 10 classes). - eval/sbatch_probe_realized_dist{,_old}.sb: validate-path run that dumps debug_mp_distribution.csv — compares the RUNTIME-realized per-op level fractions against the calibration table's intent (A4 check). - eval/sbatch_kernel_level_err.sb: kernel-level rel-err vs stoc_len under halve on qk/av/mlp shapes (catches level-grid pathologies without running the diffusion model). - sbatch_uniform48_rerun100.sb: regenerate uniform48 idx 0-99 into a fresh dir — reproducibility check. NOTE: latent noise pairs by position in each worker's index list, not by global index, so PSNR comparisons against runs with a different GPU split are only valid on the aligned worker-0 prefix. - sbatch_only1op_sc32.sb: single-op SC isolation (everything FP except one op at stoc_len 32 + halve) via per-op config JSONs; the legacy layerwise flags are ignored whenever --sc_config supplies a precision map. Co-Authored-By: Claude Fable 5 --- scripts/eval/sbatch_eval_ab_kid100.sb | 47 ++++++++++++++ scripts/eval/sbatch_eval_uniform_halve.sb | 43 +++++++++++++ scripts/eval/sbatch_kernel_level_err.sb | 61 ++++++++++++++++++ scripts/eval/sbatch_probe_realized_dist.sb | 62 +++++++++++++++++++ .../eval/sbatch_probe_realized_dist_old.sb | 51 +++++++++++++++ scripts/sbatch_only1op_sc32.sb | 58 +++++++++++++++++ scripts/sbatch_uniform48_rerun100.sb | 45 ++++++++++++++ 7 files changed, 367 insertions(+) create mode 100644 scripts/eval/sbatch_eval_ab_kid100.sb create mode 100644 scripts/eval/sbatch_eval_uniform_halve.sb create mode 100644 scripts/eval/sbatch_kernel_level_err.sb create mode 100644 scripts/eval/sbatch_probe_realized_dist.sb create mode 100644 scripts/eval/sbatch_probe_realized_dist_old.sb create mode 100644 scripts/sbatch_only1op_sc32.sb create mode 100644 scripts/sbatch_uniform48_rerun100.sb diff --git a/scripts/eval/sbatch_eval_ab_kid100.sb b/scripts/eval/sbatch_eval_ab_kid100.sb new file mode 100644 index 0000000..12e87a2 --- /dev/null +++ b/scripts/eval/sbatch_eval_ab_kid100.sb @@ -0,0 +1,47 @@ +#!/bin/bash +#SBATCH --job-name=kid_ab100 +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:30:00 +# One-off: KID on idx 0-99 for the metric A/B probe arms + uniform48 anchor. +# All three use the same indices (classes 0-9, 10 imgs each) vs the full +# ImageNet-256 ref -> absolute KID inflated by class mismatch, but the +# three-way RANKING is apples-to-apples. +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +A=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/adaptive_avg48/samples +B=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/adaptive_avg48_mse_rel/samples +U=$S/scmp_diffusion_fid_uniform_halveON_cfg15/uniform48/samples +F=$S/scmp_diffusion_fid_fp_6shard_cfg15/samples +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_ab_eval; mkdir -p "$WORK" +OUT=$WORK/kid_ab100.txt +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$WORK/eval_ab100.log") 2>&1 +echo "############ A/B KID idx0-99 (GPU) $(date) ############" + +conda activate qdit +for arm in A B U F; do + src_var=$arm; src=${!src_var} + SEL=$WORK/sel_$arm; rm -rf "$SEL"; mkdir -p "$SEL" + n=0 + for i in $(seq 0 99); do + f=$(printf '%06d.png' "$i") + [[ -f "$src/$f" ]] && ln -sf "$src/$f" "$SEL/$f" && n=$((n+1)) + done + echo "[$arm] linked $n/100 from $src" + python -u "$REPO/scripts/eval/pngs_to_npz.py" "$SEL" "$WORK/arm_$arm.npz" +done +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +python -u "$REPO/scripts/eval/kid_openai.py" "$REF" "$OUT" \ + "$WORK/arm_A.npz" "$WORK/arm_B.npz" "$WORK/arm_U.npz" "$WORK/arm_F.npz" +conda deactivate +echo "############ DONE $(date) ############" diff --git a/scripts/eval/sbatch_eval_uniform_halve.sb b/scripts/eval/sbatch_eval_uniform_halve.sb new file mode 100644 index 0000000..26e9e4e --- /dev/null +++ b/scripts/eval/sbatch_eval_uniform_halve.sb @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH --job-name=eval_uni +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:45:00 +# GPU eval for one uniform (halve) config: KID + FID/IS/sFID/Prec/Recall. +# Same methodology as sbatch_eval_mp_flop_halve.sb, repointed to the uniform dir. +# Usage: sbatch sbatch_eval_uniform_halve.sb +set -uo pipefail +AVG="${1:?usage: sbatch sbatch_eval_uniform_halve.sb }" +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +PREV=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_prev +PACKER=$PREV/imagenet256_ref/parallel_npz.py +REF=$PREV/imagenet256_ref/VIRTUAL_imagenet256_labeled.npz +BASE=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_uniform_halveON_cfg15/uniform${AVG} +AD=$BASE/samples +WORK=$BASE/_eval; mkdir -p "$WORK" +OUT=$WORK/eval_uniform${AVG}.txt +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh +exec > >(tee "$OUT") 2>&1 +echo "############ UNIFORM(halve) EVAL uniform${AVG} (GPU) $(date) ############" + +IDX=$WORK/idx.txt; NPZ=$WORK/uniform${AVG}.npz +(cd "$AD" && ls [0-9][0-9][0-9][0-9][0-9][0-9].png | awk '($1+0)<2000' | sort) > "$IDX" +SEL=$WORK/sel; rm -rf "$SEL"; mkdir -p "$SEL" +while read -r f; do [[ -n "$f" ]] && ln -sf "$AD/$f" "$SEL/$f"; done < "$IDX" +echo "=== matched idx 0-1999: $(wc -l < "$IDX") ===" + +conda activate qdit +python -u "$PACKER" "$SEL" "$NPZ" +conda deactivate + +conda activate tfeval +export TF_CPP_MIN_LOG_LEVEL=2 EVALUATOR=$PREV/Q-DiT/models/evaluations/evaluator.py +echo ""; echo "===== [1] KID vs ref (idx 0-1999) =====" +bash "$REPO/scripts/eval/kid_openai.sh" "$WORK/kid_uniform${AVG}.txt" "$NPZ" +echo ""; echo "===== [2] FID / IS / sFID / Precision / Recall vs ref =====" +python -u "$EVALUATOR" "$REF" "$NPZ" 2>&1 | grep -E "^(Inception Score|FID|sFID|Precision|Recall):" | sed "s/^/[uniform${AVG}] /" +conda deactivate +echo ""; echo "############ DONE uniform${AVG} $(date) ############" diff --git a/scripts/eval/sbatch_kernel_level_err.sb b/scripts/eval/sbatch_kernel_level_err.sb new file mode 100644 index 0000000..70b9ca0 --- /dev/null +++ b/scripts/eval/sbatch_kernel_level_err.sb @@ -0,0 +1,61 @@ +#!/bin/bash +#SBATCH --job-name=kern_lvl_err +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=24G +#SBATCH --time=00:15:00 +# Kernel-level SC error vs stoc_len under halve: is L=128 anomalously bad? +# Random matrices, per-row bipolar, qk-like and mlp-like shapes, mixed-level +# calls in one process (cache interactions included, like MP dispatch does). +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1; cd "$REPO" +exec > >(tee "$WORK/kernel_level_err.log") 2>&1 +python - << 'EOF' +import torch +from scmp_kernels.sc.matmul import sc_matmul + +torch.manual_seed(0) +dev = "cuda" + +def rel_err(out, exact): + return ((out - exact).norm() / exact.norm()).item() + +# (name, N, M, D): a is (N, D), b is (M, D), product a @ b.T +SHAPES = [ + ("qk-like (256x256, D=72)", 256, 256, 72), + ("av-like (256x72, D=256)", 256, 72, 256), + ("mlp-like (512x4608, D=1152, chunk144)", 512, 4608, 1152), +] +LEVELS = [32, 48, 64, 96, 128] + +for name, N, M, D in SHAPES: + a = torch.randn(N, D, device=dev) + b = torch.randn(M, D, device=dev) + exact = a @ b.T + chunk = 144 if D > 1024 else 0 + print(f"=== {name} ===") + # halve ON, ascending then repeated (cache warm) order + for tag, levels in (("halve asc ", LEVELS), ("halve rpt ", LEVELS)): + errs = [] + for sl in levels: + out = sc_matmul(a, b, granularity="per_row", mode="bipolar", + sc_prec=8, stoc_len=sl, group_a=1, group_b=1, + chunk_d=chunk, halve_bipolar_stoc_len=True) + errs.append(f"{sl}:{rel_err(out, exact):.4f}") + print(f" {tag}: " + " ".join(errs)) + # halve OFF reference at same lengths + 256 + errs = [] + for sl in LEVELS + [256]: + out = sc_matmul(a, b, granularity="per_row", mode="bipolar", + sc_prec=8, stoc_len=sl, group_a=1, group_b=1, + chunk_d=chunk, halve_bipolar_stoc_len=False) + errs.append(f"{sl}:{rel_err(out, exact):.4f}") + print(f" no-halve : " + " ".join(errs)) +EOF +echo "DONE $(date)" diff --git a/scripts/eval/sbatch_probe_realized_dist.sb b/scripts/eval/sbatch_probe_realized_dist.sb new file mode 100644 index 0000000..cc28e41 --- /dev/null +++ b/scripts/eval/sbatch_probe_realized_dist.sb @@ -0,0 +1,62 @@ +#!/bin/bash +#SBATCH --job-name=mp_dist_probe +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=01:00:00 +# A4 instrumented probe: run the VALIDATE path (8 imgs, dumps +# debug_mp_distribution.csv) once with the OLD 256-grid flop-halve table and +# once with the NEW 128-cap table, then print realized per-op level fractions +# side by side. 10 sampling steps -> one visit per t-bucket, enough for the +# realized-distribution comparison at ~7 min/run. +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$S/pretrained_models/DiT-XL-2-256x256.pt +OLD_CALIB=$S/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg48.json +NEW_CALIB=$S/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve_max128/calib_fix_avg48.json +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe +mkdir -p "$WORK" +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +exec > >(tee "$WORK/probe.log") 2>&1 + +run_probe () { + local tag="$1" calib="$2" levels="$3" + echo "===== probe $tag levels=$levels =====" + rm -rf "$WORK/$tag"; mkdir -p "$WORK/$tag" + python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row --sc_halve \ + --adaptive_mp --adaptive_mp_table "$calib" --mp_levels "$levels" \ + --image-size 256 --num-sampling-steps 10 --cfg-scale 1.5 \ + --seed 0 --results-dir "$WORK/$tag" --ckpt "$CKPT" + find "$WORK/$tag" -name 'debug_mp_distribution.csv' -exec cp {} "$WORK/dist_$tag.csv" \; + find "$WORK/$tag" -name 'mp_savings_summary.txt' -exec cp {} "$WORK/savings_$tag.txt" \; +} + +run_probe old "$OLD_CALIB" "256,192,128,96,64,48,32,16" +run_probe new "$NEW_CALIB" "128,96,64,48,32,16" + +python - << 'EOF' +import csv, collections +for tag in ("old", "new"): + per_op = collections.defaultdict(lambda: collections.Counter()) + tot = collections.Counter() + with open(f"/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe/dist_{tag}.csv") as f: + for row in csv.DictReader(f): + op = row["operator"] + for k, v in row.items(): + if k.startswith("sl_") and k.endswith("_count") and v: + sl = int(k.split("_")[1]); per_op[op][sl] += int(v); tot[op] += int(v) + print(f"===== REALIZED distribution ({tag} table) =====") + for op in sorted(per_op): + c = per_op[op]; n = tot[op] + avg = sum(sl * cnt for sl, cnt in c.items()) / max(n, 1) + fr = " ".join(f"{sl}:{100*cnt/n:.1f}%" for sl, cnt in sorted(c.items(), reverse=True) if cnt) + print(f" {op:12s} avg={avg:7.1f} {fr}") +EOF +echo "DONE $(date)" diff --git a/scripts/eval/sbatch_probe_realized_dist_old.sb b/scripts/eval/sbatch_probe_realized_dist_old.sb new file mode 100644 index 0000000..55ae58d --- /dev/null +++ b/scripts/eval/sbatch_probe_realized_dist_old.sb @@ -0,0 +1,51 @@ +#!/bin/bash +#SBATCH --job-name=mp_dist_old +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-gpu=4 +#SBATCH --mem=48G +#SBATCH --time=00:40:00 +# Control probe: realized distribution under the OLD 256-grid cfg0-calib table. +# Runs WITHOUT --sc_halve because the fixed kernels (correctly) refuse +# stoc_len>128 under halve — halve only changes SC noise, not the +# threshold classification, so realized level fractions remain valid. +set -uo pipefail +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +S=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$S/pretrained_models/DiT-XL-2-256x256.pt +OLD_CALIB=$S/scmp_diffusion_fid_mp_cfg15/calib_global_perrow_flop_halve/calib_fix_avg48.json +WORK=$S/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe +mkdir -p "$WORK" +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +exec > >(tee "$WORK/probe_old.log") 2>&1 + +echo "===== probe old(halve-off) levels=256,192,128,96,64,48,32,16 =====" +rm -rf "$WORK/old"; mkdir -p "$WORK/old" +python -u scripts/quant_sc_main.py \ + --wbits 8 --abits 8 --w_sym --a_sym \ + --timewise 1 --qklayerwise 1.0 --avlayerwise 1.0 --projlayerwise 1.0 --mlplayerwise 1.0 --inputprojlayerwise 1.0 \ + --sc_prec 8 --sc_fixed_level_prec --sc_qk_granularity per_row \ + --adaptive_mp --adaptive_mp_table "$OLD_CALIB" --mp_levels "256,192,128,96,64,48,32,16" \ + --image-size 256 --num-sampling-steps 10 --cfg-scale 1.5 \ + --seed 0 --results-dir "$WORK/old" --ckpt "$CKPT" +find "$WORK/old" -name 'debug_mp_distribution.csv' -exec cp {} "$WORK/dist_old.csv" \; + +python - << 'EOF' +import csv, collections +W = "/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion_fid_mp_globalpr_flop_halve128_cfg15/_dist_probe" +per_op = collections.defaultdict(collections.Counter); tot = collections.Counter() +for row in csv.DictReader(open(f"{W}/dist_old.csv")): + op = row["operator"] + for k, v in row.items(): + if k.startswith("sl_") and k.endswith("_count") and v: + per_op[op][int(k.split("_")[1])] += int(v); tot[op] += int(v) +print("===== REALIZED (old table, runtime cfg1.5, halve-off) =====") +for op in sorted(per_op): + c = per_op[op]; n = tot[op] + avg = sum(sl * cnt for sl, cnt in c.items()) / max(n, 1) + fr = " ".join(f"{sl}:{100*cnt/n:.1f}%" for sl, cnt in sorted(c.items(), reverse=True)) + print(f" {op:12s} avg={avg:7.1f} {fr}") +EOF +echo "DONE $(date)" diff --git a/scripts/sbatch_only1op_sc32.sb b/scripts/sbatch_only1op_sc32.sb new file mode 100644 index 0000000..fe2cdfa --- /dev/null +++ b/scripts/sbatch_only1op_sc32.sb @@ -0,0 +1,58 @@ +#!/bin/bash +#SBATCH --job-name=only1op32 +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=06:00:00 +# Single-op SC isolation probe (era-2 code): EVERYTHING runs FP except the ONE +# selected op-group, which runs SC at uniform stoc_len=32 + halve. Generates +# idx 0-99. The bifurcation rate vs FP directly measures that op's solo +# trajectory-branching power at a stressed length. +# v2: isolation via per-op config JSON (enabled=false for others) — the +# layerwise flags are IGNORED when --sc_config provides a precision map +# (sc_controller.py:115), which silently turned v1 into uniform32 x5. +# Usage: sbatch sbatch_only1op_sc32.sb OP in: qk av proj inproj mlp +set -uo pipefail +OP="${1:?usage: sbatch sbatch_only1op_sc32.sb }" +QK=0; AV=0; PROJ=0; INPROJ=0; MLP=0 +case "$OP" in + qk) QK=1.0 ;; + av) AV=1.0 ;; + proj) PROJ=1.0 ;; + inproj) INPROJ=1.0 ;; + mlp) MLP=1.0 ;; + *) echo "unknown OP $OP" >&2; exit 1 ;; +esac +NUM_GPUS=3 +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_only32_${OP}.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_halveON_cfg15/only1op32b_${OP} +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=100; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +echo "=== only1op SC32 op=$OP (qk=$QK av=$AV proj=$PROJ inproj=$INPROJ mlp=$MLP) job=${SLURM_JOB_ID:-?} $(date) ===" +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== only1op $OP done: $(find "$SAMPLES" -name '*.png' | wc -l)/$NUM_FID $(date) ===" diff --git a/scripts/sbatch_uniform48_rerun100.sb b/scripts/sbatch_uniform48_rerun100.sb new file mode 100644 index 0000000..0f512be --- /dev/null +++ b/scripts/sbatch_uniform48_rerun100.sb @@ -0,0 +1,45 @@ +#!/bin/bash +#SBATCH --job-name=uni48_rerun +#SBATCH --partition=gpu-rtx6000 +#SBATCH --account=nbleier_owned1 +#SBATCH --gres=gpu:3 +#SBATCH --cpus-per-gpu=2 +#SBATCH --mem-per-gpu=12G +#SBATCH --time=06:00:00 +# One-off reproducibility probe: regenerate uniform48 (halve) idx 0-99 under +# the CURRENT code state into a fresh dir, to verify the uniform48 anchor +# still reproduces (~24 dB vs FP, ~5% bifurcation) after the kernel guard / +# recent changes. Identical flags/config/seed to sbatch_uniform_halve_auto.sb. +set -uo pipefail +NUM_GPUS=3 +REPO=/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion +SCRATCH=/scratch/nbleier_owned_root/nbleier_owned1/zhkangqi +CKPT=$SCRATCH/pretrained_models/DiT-XL-2-256x256.pt +SC_JSON=$SCRATCH/scmp_diffusion_fid_halve_bitrev/configs/sc_cfg_uniform48_all.json +OUT=$SCRATCH/scmp_diffusion_fid_uniform_halveON_cfg15/uniform48_rerun +SAMPLES=$OUT/samples; IDX=$OUT/_indices; LOG=$OUT/_logs +NUM_FID=100; BALANCED=10000; BATCH=64; STEPS=50; CFG=1.5; SEED=0; NUM_CLASSES=1000 +mkdir -p "$SAMPLES" "$IDX" "$LOG" +[[ -f "$SC_JSON" ]] || { echo "ERROR: missing $SC_JSON" >&2; exit 1; } +source /home/zhkangqi/miniconda3/etc/profile.d/conda.sh; conda activate qdit +export SC_OWEN_MODE=bitrev PYTHONUNBUFFERED=1 OMP_NUM_THREADS=2; cd "$REPO" +echo "=== uniform48 RERUN probe: target=$NUM_FID job=${SLURM_JOB_ID:-?} $(date) ===" +python -u scripts/_plan_missing_indices.py "$SAMPLES" "$NUM_FID" "$NUM_GPUS" "$IDX" $((BALANCED/NUM_CLASSES)) +pids=() +for ((g=0; g "$LOG/gpu_${g}.log" 2>&1 & + pids+=($!) +done +for p in "${pids[@]}"; do wait "$p" || true; done +echo "=== uniform48 RERUN done: $(find "$SAMPLES" -name '*.png' | wc -l)/$NUM_FID $(date) ==="