remdm-planner-minihack / scripts /profile_ablation.py

Refresh artefacts and code for the second review release

5c30113 verified 2 months ago

10.2 kB

"""Profile one RL fine-tuning ablation iteration by component.
Runs ``run_ablation`` unmodified and attributes wall-clock time to the
pieces that matter for the speed pass: environment construction, env
stepping, GPU sampling, window extraction, the host-to-device transfer,
the gradient step, each diagnostic, and evaluation.
Everything is measured by monkey-patching the callables ``training.py``
already uses, so the training loop itself is never edited and the
measured run computes exactly what a real run computes.
Usage:
    python scripts/profile_ablation.py --checkpoint PATH \
        --ablation baseline_rl --max-iter 51
Without ``--config`` this profiles ``configs/defaults.yaml`` merged with
``experiments/rl_finetuning/configs/ablations_final_minihack_gpu_24gb.yaml``.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
from collections import defaultdict
from pathlib import Path

import torch

_PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from experiments.rl_finetuning.ablations import training as T  # noqa: E402
from experiments.rl_finetuning.ablations.registry import REGISTRY  # noqa: E402
from experiments.rl_finetuning.run_ablations import (  # noqa: E402
    _load_yaml,
    _merge_to_namespace,
)
from src.envs import minihack_env  # noqa: E402
from src.models.denoiser import ModelEMA  # noqa: E402
from src.planners.inference import Evaluator  # noqa: E402

logger = logging.getLogger("profile_ablation")


class Acc:
    """Accumulating wall-clock timer with a call counter."""

    def __init__(self) -> None:
        self.total = 0.0
        self.calls = 0

    def add(self, dt: float) -> None:
        """Record one timed call of duration *dt* seconds."""
        self.total += dt
        self.calls += 1

    def reset(self) -> None:
        """Zero the accumulator."""
        self.total = 0.0
        self.calls = 0


ACC: dict[str, Acc] = defaultdict(Acc)
WINDOWS: list[int] = []


def _sync() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def _wrap(fn, key: str, sync: bool = False):
    """Return *fn* wrapped so its wall time accumulates under *key*."""

    def inner(*a, **kw):
        if sync:
            _sync()
        t0 = time.perf_counter()
        try:
            return fn(*a, **kw)
        finally:
            if sync:
                _sync()
            ACC[key].add(time.perf_counter() - t0)

    return inner


def install_patches() -> None:
    """Monkey-patch the callables whose cost we want attributed."""
    # -- Collection internals -------------------------------------------
    # ``collect/env_acquire`` is the whole cost of getting a usable env,
    # whether that is a fresh construction or a pool hit;
    # ``collect/env_init`` is the construction alone, so the two diverge
    # once the pool starts serving.
    if hasattr(T, "acquire_env"):
        T.acquire_env = _wrap(T.acquire_env, "collect/env_acquire")
    else:
        T.make_env = _wrap(T.make_env, "collect/env_acquire")
    minihack_env.AdvancedObservationEnv.step = _wrap(
        minihack_env.AdvancedObservationEnv.step, "collect/env_step"
    )
    minihack_env.AdvancedObservationEnv.reset = _wrap(
        minihack_env.AdvancedObservationEnv.reset, "collect/env_reset"
    )
    minihack_env.AdvancedObservationEnv.close = _wrap(
        minihack_env.AdvancedObservationEnv.close, "collect/env_close"
    )
    minihack_env.AdvancedObservationEnv.__init__ = _wrap(
        minihack_env.AdvancedObservationEnv.__init__, "collect/env_init"
    )
    T.remdm_sample = _wrap(T.remdm_sample, "collect/gpu_sample", sync=True)
    T._extract_windows = _wrap(T._extract_windows, "collect/extract_windows")

    # -- Window count actually available to the gradient step ------------
    # ``compute_advantages`` is called once per iteration on every window
    # collected after the ablation's filters, so its input length is the
    # number the ``local_obs[:batch_size]`` slice draws from.
    _adv = T.compute_advantages

    def _adv_counting(returns, *a, **kw):
        WINDOWS.append(int(returns.shape[0]))
        return _adv(returns, *a, **kw)

    T.compute_advantages = _adv_counting

    # -- Model plumbing --------------------------------------------------
    ModelEMA.make_eval_model = _wrap(
        ModelEMA.make_eval_model, "train/make_eval_model", sync=True
    )

    # -- Diagnostics -----------------------------------------------------
    T.compute_grad_alignment = _wrap(
        T.compute_grad_alignment, "diag/grad_alignment", sync=True
    )
    T.compute_per_layer_grad_norms = _wrap(
        T.compute_per_layer_grad_norms, "diag/per_layer_norms_only", sync=True
    )
    T.compute_repr_drift = _wrap(T.compute_repr_drift, "diag/repr_drift", sync=True)
    T.compute_cka = _wrap(T.compute_cka, "diag/cka", sync=True)
    T.compute_t_analysis = _wrap(T.compute_t_analysis, "diag/t_analysis", sync=True)
    Evaluator.evaluate = _wrap(Evaluator.evaluate, "eval/evaluate", sync=True)


ITER_ROWS: list[dict] = []


def install_metric_capture() -> None:
    """Capture each iteration's metric dict plus the timer deltas.
    ``training.py`` stops its own ``speed/iter_time_sec`` clock straight
    after the gradient step, before the diagnostics and the eval, so it
    understates a diagnostic iteration badly. ``prof/wall_iter_sec`` is
    the gap between consecutive log calls, which is the real thing.
    """
    prev_total: dict[str, float] = {}
    prev_calls: dict[str, int] = {}
    last = [time.perf_counter()]
    seen_windows = [0]

    def _log(metrics: dict, step: int) -> None:
        now = time.perf_counter()
        row = dict(metrics)
        row["prof/wall_iter_sec"] = now - last[0]
        last[0] = now
        # First new advantage call of the iteration is the online batch;
        # mixed_replay adds a second call for the buffer sample.
        if len(WINDOWS) > seen_windows[0]:
            row["prof/windows_collected"] = WINDOWS[seen_windows[0]]
            seen_windows[0] = len(WINDOWS)
        for key, acc in ACC.items():
            row[f"prof/{key}"] = acc.total - prev_total.get(key, 0.0)
            row[f"prof/{key}_calls"] = acc.calls - prev_calls.get(key, 0)
            prev_total[key] = acc.total
            prev_calls[key] = acc.calls
        ITER_ROWS.append(row)

    T._wandb_log = _log


def main() -> None:
    """Entry point."""
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--config", default=str(_PROJECT_ROOT / "configs/defaults.yaml"))
    p.add_argument(
        "--ablations-config",
        default=str(
            _PROJECT_ROOT
            / "experiments/rl_finetuning/configs/ablations_final_minihack_gpu_24gb.yaml"
        ),
    )
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--ablation", default="baseline_rl")
    p.add_argument("--max-iter", type=int, default=51)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--device", default="cuda")
    p.add_argument("--out", default=None, help="Write per-iteration JSON here.")
    p.add_argument(
        "--override",
        action="append",
        default=[],
        help="Config override key=value (parsed as JSON, then as str).",
    )
    args = p.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    overrides: dict = {"max_iter": args.max_iter}
    for item in args.override:
        key, _, raw = item.partition("=")
        try:
            overrides[key] = json.loads(raw)
        except json.JSONDecodeError:
            overrides[key] = raw

    cfg = _merge_to_namespace(
        _load_yaml(args.config),
        _load_yaml(args.ablations_config),
        overrides,
    )
    device = torch.device(args.device)
    if device.type == "cuda":
        torch.set_float32_matmul_precision("high")

    logger.info("TMPDIR=%s", os.environ.get("TMPDIR", "<unset>"))
    logger.info("config=%s", args.config)
    logger.info("ablations_config=%s", args.ablations_config)
    logger.info("checkpoint=%s", args.checkpoint)
    logger.info(
        "n_embd=%s n_head=%s n_layer=%s batch_size=%s episodes_per_iter=%s",
        cfg.n_embd,
        cfg.n_head,
        cfg.n_layer,
        cfg.batch_size,
        cfg.episodes_per_iter,
    )

    install_patches()
    install_metric_capture()

    spec = REGISTRY[args.ablation]
    t0 = time.perf_counter()
    _, final_score, _, _ = T.run_ablation(
        spec, cfg, args.checkpoint, device, seed=args.seed
    )
    wall = time.perf_counter() - t0

    print("\n" + "=" * 78)
    print(f"ablation={args.ablation} iters={args.max_iter} wall={wall:.1f}s")
    print(f"final_score={final_score:.4f}")
    print("=" * 78)
    print(f"{'component':38s} {'total_s':>10s} {'calls':>9s} {'ms/call':>10s}")
    for key in sorted(ACC):
        acc = ACC[key]
        per = 1000.0 * acc.total / acc.calls if acc.calls else 0.0
        print(f"{key:38s} {acc.total:10.2f} {acc.calls:9d} {per:10.3f}")

    if args.out:
        Path(args.out).write_text(
            json.dumps(
                {
                    "ablation": args.ablation,
                    "max_iter": args.max_iter,
                    "seed": args.seed,
                    "wall_sec": wall,
                    "final_score": final_score,
                    "config": {
                        "n_embd": cfg.n_embd,
                        "n_head": cfg.n_head,
                        "n_layer": cfg.n_layer,
                        "batch_size": cfg.batch_size,
                        "episodes_per_iter": cfg.episodes_per_iter,
                    },
                    "totals": {
                        k: {"total_sec": v.total, "calls": v.calls}
                        for k, v in ACC.items()
                    },
                    "iters": ITER_ROWS,
                },
                indent=1,
            )
        )
        print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()

	"""Profile one RL fine-tuning ablation iteration by component.

	Runs ``run_ablation`` unmodified and attributes wall-clock time to the
	pieces that matter for the speed pass: environment construction, env
	stepping, GPU sampling, window extraction, the host-to-device transfer,
	the gradient step, each diagnostic, and evaluation.

	Everything is measured by monkey-patching the callables ``training.py``
	already uses, so the training loop itself is never edited and the
	measured run computes exactly what a real run computes.

	Usage:
	python scripts/profile_ablation.py --checkpoint PATH \
	--ablation baseline_rl --max-iter 51

	Without ``--config`` this profiles ``configs/defaults.yaml`` merged with
	``experiments/rl_finetuning/configs/ablations_final_minihack_gpu_24gb.yaml``.
	"""

	from __future__ import annotations

	import argparse
	import json
	import logging
	import os
	import sys
	import time
	from collections import defaultdict
	from pathlib import Path

	import torch

	_PROJECT_ROOT = Path(__file__).resolve().parents[1]
	if str(_PROJECT_ROOT) not in sys.path:
	sys.path.insert(0, str(_PROJECT_ROOT))

	from experiments.rl_finetuning.ablations import training as T # noqa: E402
	from experiments.rl_finetuning.ablations.registry import REGISTRY # noqa: E402
	from experiments.rl_finetuning.run_ablations import ( # noqa: E402
	_load_yaml,
	_merge_to_namespace,
	)
	from src.envs import minihack_env # noqa: E402
	from src.models.denoiser import ModelEMA # noqa: E402
	from src.planners.inference import Evaluator # noqa: E402

	logger = logging.getLogger("profile_ablation")


	class Acc:
	"""Accumulating wall-clock timer with a call counter."""

	def __init__(self) -> None:
	self.total = 0.0
	self.calls = 0

	def add(self, dt: float) -> None:
	"""Record one timed call of duration dt seconds."""
	self.total += dt
	self.calls += 1

	def reset(self) -> None:
	"""Zero the accumulator."""
	self.total = 0.0
	self.calls = 0


	ACC: dict[str, Acc] = defaultdict(Acc)
	WINDOWS: list[int] = []


	def _sync() -> None:
	if torch.cuda.is_available():
	torch.cuda.synchronize()


	def _wrap(fn, key: str, sync: bool = False):
	"""Return fn wrapped so its wall time accumulates under key."""

	def inner(a, *kw):
	if sync:
	_sync()
	t0 = time.perf_counter()
	try:
	return fn(a, *kw)
	finally:
	if sync:
	_sync()
	ACC[key].add(time.perf_counter() - t0)

	return inner


	def install_patches() -> None:
	"""Monkey-patch the callables whose cost we want attributed."""
	# -- Collection internals -------------------------------------------
	# ``collect/env_acquire`` is the whole cost of getting a usable env,
	# whether that is a fresh construction or a pool hit;
	# ``collect/env_init`` is the construction alone, so the two diverge
	# once the pool starts serving.
	if hasattr(T, "acquire_env"):
	T.acquire_env = _wrap(T.acquire_env, "collect/env_acquire")
	else:
	T.make_env = _wrap(T.make_env, "collect/env_acquire")
	minihack_env.AdvancedObservationEnv.step = _wrap(
	minihack_env.AdvancedObservationEnv.step, "collect/env_step"
	)
	minihack_env.AdvancedObservationEnv.reset = _wrap(
	minihack_env.AdvancedObservationEnv.reset, "collect/env_reset"
	)
	minihack_env.AdvancedObservationEnv.close = _wrap(
	minihack_env.AdvancedObservationEnv.close, "collect/env_close"
	)
	minihack_env.AdvancedObservationEnv.__init__ = _wrap(
	minihack_env.AdvancedObservationEnv.__init__, "collect/env_init"
	)
	T.remdm_sample = _wrap(T.remdm_sample, "collect/gpu_sample", sync=True)
	T._extract_windows = _wrap(T._extract_windows, "collect/extract_windows")

	# -- Window count actually available to the gradient step ------------
	# ``compute_advantages`` is called once per iteration on every window
	# collected after the ablation's filters, so its input length is the
	# number the ``local_obs[:batch_size]`` slice draws from.
	_adv = T.compute_advantages

	def _adv_counting(returns, a, *kw):
	WINDOWS.append(int(returns.shape[0]))
	return _adv(returns, a, *kw)

	T.compute_advantages = _adv_counting

	# -- Model plumbing --------------------------------------------------
	ModelEMA.make_eval_model = _wrap(
	ModelEMA.make_eval_model, "train/make_eval_model", sync=True
	)

	# -- Diagnostics -----------------------------------------------------
	T.compute_grad_alignment = _wrap(
	T.compute_grad_alignment, "diag/grad_alignment", sync=True
	)
	T.compute_per_layer_grad_norms = _wrap(
	T.compute_per_layer_grad_norms, "diag/per_layer_norms_only", sync=True
	)
	T.compute_repr_drift = _wrap(T.compute_repr_drift, "diag/repr_drift", sync=True)
	T.compute_cka = _wrap(T.compute_cka, "diag/cka", sync=True)
	T.compute_t_analysis = _wrap(T.compute_t_analysis, "diag/t_analysis", sync=True)
	Evaluator.evaluate = _wrap(Evaluator.evaluate, "eval/evaluate", sync=True)


	ITER_ROWS: list[dict] = []


	def install_metric_capture() -> None:
	"""Capture each iteration's metric dict plus the timer deltas.

	``training.py`` stops its own ``speed/iter_time_sec`` clock straight
	after the gradient step, before the diagnostics and the eval, so it
	understates a diagnostic iteration badly. ``prof/wall_iter_sec`` is
	the gap between consecutive log calls, which is the real thing.
	"""
	prev_total: dict[str, float] = {}
	prev_calls: dict[str, int] = {}
	last = [time.perf_counter()]
	seen_windows = [0]

	def _log(metrics: dict, step: int) -> None:
	now = time.perf_counter()
	row = dict(metrics)
	row["prof/wall_iter_sec"] = now - last[0]
	last[0] = now
	# First new advantage call of the iteration is the online batch;
	# mixed_replay adds a second call for the buffer sample.
	if len(WINDOWS) > seen_windows[0]:
	row["prof/windows_collected"] = WINDOWS[seen_windows[0]]
	seen_windows[0] = len(WINDOWS)
	for key, acc in ACC.items():
	row[f"prof/{key}"] = acc.total - prev_total.get(key, 0.0)
	row[f"prof/{key}_calls"] = acc.calls - prev_calls.get(key, 0)
	prev_total[key] = acc.total
	prev_calls[key] = acc.calls
	ITER_ROWS.append(row)

	T._wandb_log = _log


	def main() -> None:
	"""Entry point."""
	p = argparse.ArgumentParser(description=__doc__)
	p.add_argument("--config", default=str(_PROJECT_ROOT / "configs/defaults.yaml"))
	p.add_argument(
	"--ablations-config",
	default=str(
	_PROJECT_ROOT
	/ "experiments/rl_finetuning/configs/ablations_final_minihack_gpu_24gb.yaml"
	),
	)
	p.add_argument("--checkpoint", required=True)
	p.add_argument("--ablation", default="baseline_rl")
	p.add_argument("--max-iter", type=int, default=51)
	p.add_argument("--seed", type=int, default=0)
	p.add_argument("--device", default="cuda")
	p.add_argument("--out", default=None, help="Write per-iteration JSON here.")
	p.add_argument(
	"--override",
	action="append",
	default=[],
	help="Config override key=value (parsed as JSON, then as str).",
	)
	args = p.parse_args()

	logging.basicConfig(
	level=logging.INFO,
	format="%(asctime)s %(levelname)s %(message)s",
	)

	overrides: dict = {"max_iter": args.max_iter}
	for item in args.override:
	key, _, raw = item.partition("=")
	try:
	overrides[key] = json.loads(raw)
	except json.JSONDecodeError:
	overrides[key] = raw

	cfg = _merge_to_namespace(
	_load_yaml(args.config),
	_load_yaml(args.ablations_config),
	overrides,
	)
	device = torch.device(args.device)
	if device.type == "cuda":
	torch.set_float32_matmul_precision("high")

	logger.info("TMPDIR=%s", os.environ.get("TMPDIR", "<unset>"))
	logger.info("config=%s", args.config)
	logger.info("ablations_config=%s", args.ablations_config)
	logger.info("checkpoint=%s", args.checkpoint)
	logger.info(
	"n_embd=%s n_head=%s n_layer=%s batch_size=%s episodes_per_iter=%s",
	cfg.n_embd,
	cfg.n_head,
	cfg.n_layer,
	cfg.batch_size,
	cfg.episodes_per_iter,
	)

	install_patches()
	install_metric_capture()

	spec = REGISTRY[args.ablation]
	t0 = time.perf_counter()
	_, final_score, _, _ = T.run_ablation(
	spec, cfg, args.checkpoint, device, seed=args.seed
	)
	wall = time.perf_counter() - t0

	print("\n" + "=" * 78)
	print(f"ablation={args.ablation} iters={args.max_iter} wall={wall:.1f}s")
	print(f"final_score={final_score:.4f}")
	print("=" * 78)
	print(f"{'component':38s} {'total_s':>10s} {'calls':>9s} {'ms/call':>10s}")
	for key in sorted(ACC):
	acc = ACC[key]
	per = 1000.0 * acc.total / acc.calls if acc.calls else 0.0
	print(f"{key:38s} {acc.total:10.2f} {acc.calls:9d} {per:10.3f}")

	if args.out:
	Path(args.out).write_text(
	json.dumps(
	{
	"ablation": args.ablation,
	"max_iter": args.max_iter,
	"seed": args.seed,
	"wall_sec": wall,
	"final_score": final_score,
	"config": {
	"n_embd": cfg.n_embd,
	"n_head": cfg.n_head,
	"n_layer": cfg.n_layer,
	"batch_size": cfg.batch_size,
	"episodes_per_iter": cfg.episodes_per_iter,
	},
	"totals": {
	k: {"total_sec": v.total, "calls": v.calls}
	for k, v in ACC.items()
	},
	"iters": ITER_ROWS,
	},
	indent=1,
	)
	)
	print(f"\nwrote {args.out}")


	if __name__ == "__main__":
	main()