admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-08 18:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(profiling): add weekly model profiling

Browse Source
This commit is contained in:
Pepijn
2026-04-15 22:31:44 +02:00
parent bd74f6733d
commit 1a2aec1b04
7 changed files with 1307 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/configs/train.py
+20
View File
@@ -56,6 +56,16 @@ class TrainPipelineConfig(HubMixin):
# Number of workers for the dataloader.
num_workers: int = 4
batch_size: int = 8
profile_mode: str = "off"
profile_wait_steps: int = 1
profile_warmup_steps: int = 2
profile_active_steps: int = 6
profile_repeat: int = 1
profile_output_dir: Path | None = None
profile_record_shapes: bool = True
profile_with_memory: bool = True
profile_with_flops: bool = True
profile_with_stack: bool = False
steps: int = 100_000
eval_freq: int = 20_000
log_freq: int = 200
@@ -128,9 +138,19 @@ class TrainPipelineConfig(HubMixin):
now = dt.datetime.now()
train_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
self.output_dir = Path("outputs/train") / train_dir
if self.profile_mode != "off" and self.profile_output_dir is None:
self.profile_output_dir = self.output_dir / "profiling"
if isinstance(self.dataset.repo_id, list):
raise NotImplementedError("LeRobotMultiDataset is not currently implemented.")
if self.profile_mode not in {"off", "summary", "trace"}:
raise ValueError(
f"`profile_mode` must be one of 'off', 'summary', or 'trace', got {self.profile_mode}."
)
if self.profile_wait_steps < 0 or self.profile_warmup_steps < 0 or self.profile_active_steps < 0:
raise ValueError("Profiler schedule steps must be non-negative.")
if self.profile_repeat <= 0:
raise ValueError("`profile_repeat` must be strictly positive.")
if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None):
raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.")
src/lerobot/scripts/lerobot_train.py
+225 -115
View File
@@ -22,6 +22,7 @@ import dataclasses
import logging
import time
from contextlib import nullcontext
from pathlib import Path
from pprint import pformat
from typing import TYPE_CHECKING, Any
@@ -49,6 +50,14 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.profiling_utils import (
StepTimingCollector,
ensure_dir,
make_torch_profiler,
run_with_cprofile,
write_deterministic_forward_artifacts,
write_torch_profiler_outputs,
)
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
cycle,
@@ -71,6 +80,7 @@ def update_policy(
lr_scheduler=None,
lock=None,
rabc_weights_provider=None,
timing_collector: StepTimingCollector | None = None,
) -> tuple[MetricsTracker, dict]:
"""
Performs a single training step to update the policy's weights.
@@ -104,6 +114,7 @@ def update_policy(
rabc_batch_weights, rabc_batch_stats = rabc_weights_provider.compute_batch_weights(batch)
# Let accelerator handle mixed precision
forward_start = time.perf_counter()
with accelerator.autocast():
# Use per-sample loss when RA-BC is enabled for proper weighting
if rabc_batch_weights is not None:
@@ -122,11 +133,15 @@ def update_policy(
loss, output_dict = policy.forward(batch)
# TODO(rcadene): policy.unnormalize_outputs(out_dict)
forward_s = time.perf_counter() - forward_start
# Use accelerator's backward method
backward_start = time.perf_counter()
accelerator.backward(loss)
backward_s = time.perf_counter() - backward_start
# Clip gradients if specified
optimizer_start = time.perf_counter()
if grad_clip_norm > 0:
grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
else:
@@ -147,11 +162,19 @@ def update_policy(
# Update internal buffers if policy has update method
if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
optimizer_s = time.perf_counter() - optimizer_start
train_metrics.loss = loss.item()
train_metrics.grad_norm = grad_norm.item()
train_metrics.lr = optimizer.param_groups[0]["lr"]
train_metrics.update_s = time.perf_counter() - start_time
if timing_collector is not None:
timing_collector.record(
forward_s=forward_s,
backward_s=backward_s,
optimizer_s=optimizer_s,
total_update_s=train_metrics.update_s,
)
return train_metrics, output_dict
@@ -206,6 +229,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
logging.info(pformat(cfg.to_dict()))
profiling_enabled = cfg.profile_mode != "off"
profile_output_dir = None
cprofile_dir = None
if profiling_enabled and is_main_process and cfg.profile_output_dir is not None:
profile_output_dir = ensure_dir(Path(cfg.profile_output_dir))
cprofile_dir = ensure_dir(profile_output_dir / "cprofile")
logging.info("Profiling enabled. Artifacts will be written to %s", profile_output_dir)
# Initialize wandb only on main process
if cfg.wandb.enable and cfg.wandb.project and is_main_process:
wandb_logger = WandBLogger(cfg)
@@ -229,7 +260,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# Dataset loading synchronization: main process downloads first to avoid race conditions
if is_main_process:
logging.info("Creating dataset")
dataset = make_dataset(cfg)
if cprofile_dir is not None:
dataset = run_with_cprofile("dataset_setup", cprofile_dir, make_dataset, cfg)
else:
dataset = make_dataset(cfg)
accelerator.wait_for_everyone()
@@ -247,11 +281,21 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
logging.info("Creating policy")
policy = make_policy(
cfg=cfg.policy,
ds_meta=dataset.meta,
rename_map=cfg.rename_map,
)
if is_main_process and cprofile_dir is not None:
policy = run_with_cprofile(
"policy_setup",
cprofile_dir,
make_policy,
cfg=cfg.policy,
ds_meta=dataset.meta,
rename_map=cfg.rename_map,
)
else:
policy = make_policy(
cfg=cfg.policy,
ds_meta=dataset.meta,
rename_map=cfg.rename_map,
)
if cfg.peft is not None:
logging.info("Using PEFT! Wrapping model.")
@@ -305,16 +349,47 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
},
}
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=cfg.policy,
pretrained_path=processor_pretrained_path,
**processor_kwargs,
**postprocessor_kwargs,
)
if is_main_process and cprofile_dir is not None:
preprocessor, postprocessor = run_with_cprofile(
"processor_setup",
cprofile_dir,
make_pre_post_processors,
policy_cfg=cfg.policy,
pretrained_path=processor_pretrained_path,
**processor_kwargs,
**postprocessor_kwargs,
)
else:
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=cfg.policy,
pretrained_path=processor_pretrained_path,
**processor_kwargs,
**postprocessor_kwargs,
)
if is_main_process:
logging.info("Creating optimizer and scheduler")
optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
if is_main_process and cprofile_dir is not None:
optimizer, lr_scheduler = run_with_cprofile(
"optimizer_setup",
cprofile_dir,
make_optimizer_and_scheduler,
cfg,
policy,
)
else:
optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
if profiling_enabled and is_main_process and profile_output_dir is not None:
logging.info("Recording deterministic forward-pass artifacts")
write_deterministic_forward_artifacts(
policy=policy,
dataset=dataset,
batch_size=cfg.batch_size,
preprocessor=preprocessor,
output_dir=profile_output_dir,
device_type=device.type,
)
# Load precomputed SARM progress for RA-BC if enabled
# Generate progress using: src/lerobot/policies/sarm/compute_rabc_weights.py
@@ -429,124 +504,159 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
logging.info(
f"Start offline training on a fixed dataset, with effective batch size: {effective_batch_size}"
)
timing_collector = StepTimingCollector() if profiling_enabled and is_main_process else None
profiler = None
profiler_context = nullcontext()
if profiling_enabled and is_main_process and profile_output_dir is not None:
if device.type == "cuda":
torch.cuda.reset_peak_memory_stats(device)
profiler = make_torch_profiler(cfg, profile_output_dir, device.type)
profiler_context = profiler
for _ in range(step, cfg.steps):
start_time = time.perf_counter()
batch = next(dl_iter)
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
with profiler_context:
for _ in range(step, cfg.steps):
start_time = time.perf_counter()
batch = next(dl_iter)
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
train_tracker, output_dict = update_policy(
train_tracker,
policy,
batch,
optimizer,
cfg.optimizer.grad_clip_norm,
accelerator=accelerator,
lr_scheduler=lr_scheduler,
rabc_weights_provider=rabc_weights,
)
train_tracker, output_dict = update_policy(
train_tracker,
policy,
batch,
optimizer,
cfg.optimizer.grad_clip_norm,
accelerator=accelerator,
lr_scheduler=lr_scheduler,
rabc_weights_provider=rabc_weights,
timing_collector=timing_collector,
)
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
# increment `step` here.
step += 1
if is_main_process:
progbar.update(1)
train_tracker.step()
is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 and is_main_process
is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
if is_log_step:
logging.info(train_tracker)
if wandb_logger:
wandb_log_dict = train_tracker.to_dict()
if output_dict:
wandb_log_dict.update(output_dict)
# Log RA-BC statistics if enabled
if rabc_weights is not None:
rabc_stats = rabc_weights.get_stats()
wandb_log_dict.update(
{
"rabc_delta_mean": rabc_stats["delta_mean"],
"rabc_delta_std": rabc_stats["delta_std"],
"rabc_num_frames": rabc_stats["num_frames"],
}
)
wandb_logger.log_dict(wandb_log_dict, step)
train_tracker.reset_averages()
if cfg.save_checkpoint and is_saving_step:
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
# increment `step` here.
step += 1
if is_main_process:
logging.info(f"Checkpoint policy after step {step}")
checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
save_checkpoint(
checkpoint_dir=checkpoint_dir,
step=step,
cfg=cfg,
policy=accelerator.unwrap_model(policy),
optimizer=optimizer,
scheduler=lr_scheduler,
preprocessor=preprocessor,
postprocessor=postprocessor,
)
update_last_checkpoint(checkpoint_dir)
progbar.update(1)
if timing_collector is not None:
timing_collector.record_dataloading(train_tracker.dataloading_s.val)
if device.type == "cuda":
timing_collector.record_memory(
step=step,
allocated_bytes=torch.cuda.memory_allocated(device),
reserved_bytes=torch.cuda.memory_reserved(device),
)
train_tracker.step()
if profiler is not None:
profiler.step()
is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 and is_main_process
is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
if is_log_step:
logging.info(train_tracker)
if wandb_logger:
wandb_logger.log_policy(checkpoint_dir)
wandb_log_dict = train_tracker.to_dict()
if output_dict:
wandb_log_dict.update(output_dict)
# Log RA-BC statistics if enabled
if rabc_weights is not None:
rabc_stats = rabc_weights.get_stats()
wandb_log_dict.update(
{
"rabc_delta_mean": rabc_stats["delta_mean"],
"rabc_delta_std": rabc_stats["delta_std"],
"rabc_num_frames": rabc_stats["num_frames"],
}
)
wandb_logger.log_dict(wandb_log_dict, step)
train_tracker.reset_averages()
accelerator.wait_for_everyone()
if cfg.env and is_eval_step:
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
with torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
if cfg.save_checkpoint and is_saving_step:
if is_main_process:
logging.info(f"Checkpoint policy after step {step}")
checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
save_checkpoint(
checkpoint_dir=checkpoint_dir,
step=step,
cfg=cfg,
policy=accelerator.unwrap_model(policy),
env_preprocessor=env_preprocessor,
env_postprocessor=env_postprocessor,
optimizer=optimizer,
scheduler=lr_scheduler,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=cfg.eval.n_episodes,
videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
max_episodes_rendered=4,
start_seed=cfg.seed,
max_parallel_tasks=cfg.env.max_parallel_tasks,
)
# overall metrics (suite-agnostic)
aggregated = eval_info["overall"]
update_last_checkpoint(checkpoint_dir)
if wandb_logger:
wandb_logger.log_policy(checkpoint_dir)
# optional: per-suite logging
for suite, suite_info in eval_info.items():
logging.info("Suite %s aggregated: %s", suite, suite_info)
accelerator.wait_for_everyone()
# meters/tracker
eval_metrics = {
"avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
"pc_success": AverageMeter("success", ":.1f"),
"eval_s": AverageMeter("eval_s", ":.3f"),
}
eval_tracker = MetricsTracker(
cfg.batch_size,
dataset.num_frames,
dataset.num_episodes,
eval_metrics,
initial_step=step,
accelerator=accelerator,
)
eval_tracker.eval_s = aggregated.pop("eval_s")
eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward")
eval_tracker.pc_success = aggregated.pop("pc_success")
if wandb_logger:
wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
if cfg.env and is_eval_step:
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
with torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy),
env_preprocessor=env_preprocessor,
env_postprocessor=env_postprocessor,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=cfg.eval.n_episodes,
videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
max_episodes_rendered=4,
start_seed=cfg.seed,
max_parallel_tasks=cfg.env.max_parallel_tasks,
)
# overall metrics (suite-agnostic)
aggregated = eval_info["overall"]
accelerator.wait_for_everyone()
# optional: per-suite logging
for suite, suite_info in eval_info.items():
logging.info("Suite %s aggregated: %s", suite, suite_info)
# meters/tracker
eval_metrics = {
"avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
"pc_success": AverageMeter("success", ":.1f"),
"eval_s": AverageMeter("eval_s", ":.3f"),
}
eval_tracker = MetricsTracker(
cfg.batch_size,
dataset.num_frames,
dataset.num_episodes,
eval_metrics,
initial_step=step,
accelerator=accelerator,
)
eval_tracker.eval_s = aggregated.pop("eval_s")
eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward")
eval_tracker.pc_success = aggregated.pop("pc_success")
if wandb_logger:
wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
accelerator.wait_for_everyone()
if is_main_process:
progbar.close()
if timing_collector is not None and profile_output_dir is not None:
extra_profile_metrics = {
"profile_mode": cfg.profile_mode,
"peak_memory_allocated_bytes": (
torch.cuda.max_memory_allocated(device) if device.type == "cuda" else None
),
"peak_memory_reserved_bytes": (
torch.cuda.max_memory_reserved(device) if device.type == "cuda" else None
),
}
timing_collector.write_json(
profile_output_dir / "step_timing_summary.json", extra=extra_profile_metrics
)
if profiler is not None and profile_output_dir is not None:
write_torch_profiler_outputs(profiler, profile_output_dir, device_type=device.type)
if eval_env:
close_envs(eval_env)
src/lerobot/utils/profiling_utils.py
+297
View File
@@ -0,0 +1,297 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import cProfile
import hashlib
import io
import json
import pstats
import statistics
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import torch
from torch.utils.data._utils.collate import default_collate
def ensure_dir(path: Path) -> Path:
path.mkdir(parents=True, exist_ok=True)
return path
def render_cprofile_summary(
profile: cProfile.Profile, *, sort_by: str = "cumulative", limit: int = 40
) -> str:
output = io.StringIO()
stats = pstats.Stats(profile, stream=output).strip_dirs().sort_stats(sort_by)
stats.print_stats(limit)
return output.getvalue()
def write_profiler_table(
profiler: Any,
output_path: Path,
*,
sort_by: str,
row_limit: int = 40,
) -> None:
try:
table = profiler.key_averages().table(sort_by=sort_by, row_limit=row_limit)
except Exception:
return
output_path.write_text(table)
def make_torch_profiler(cfg: Any, output_dir: Path, device_type: str) -> Any:
activities = [torch.profiler.ProfilerActivity.CPU]
if device_type == "cuda":
activities.append(torch.profiler.ProfilerActivity.CUDA)
trace_dir = ensure_dir(output_dir / "torch_traces")
def _trace_ready(profiler: Any) -> None:
if cfg.profile_mode != "trace":
return
profiler.export_chrome_trace(str(trace_dir / f"trace_step_{profiler.step_num}.json"))
return torch.profiler.profile(
activities=activities,
schedule=torch.profiler.schedule(
wait=cfg.profile_wait_steps,
warmup=cfg.profile_warmup_steps,
active=cfg.profile_active_steps,
repeat=cfg.profile_repeat,
),
on_trace_ready=_trace_ready,
record_shapes=cfg.profile_record_shapes,
profile_memory=cfg.profile_with_memory,
with_flops=cfg.profile_with_flops,
with_stack=cfg.profile_with_stack,
)
def write_torch_profiler_outputs(
profiler: Any,
output_dir: Path,
*,
device_type: str,
) -> None:
tables_dir = ensure_dir(output_dir / "torch_tables")
write_profiler_table(profiler, tables_dir / "cpu_time_total.txt", sort_by="cpu_time_total")
if device_type == "cuda":
write_profiler_table(profiler, tables_dir / "cuda_time_total.txt", sort_by="self_cuda_time_total")
write_profiler_table(profiler, tables_dir / "cuda_memory.txt", sort_by="self_cuda_memory_usage")
write_profiler_table(profiler, tables_dir / "cpu_memory.txt", sort_by="self_cpu_memory_usage")
write_profiler_table(profiler, tables_dir / "flops.txt", sort_by="flops")
def run_with_cprofile[T](
label: str,
output_dir: Path,
fn: Callable[..., T],
*args: Any,
sort_by: str = "cumulative",
limit: int = 40,
**kwargs: Any,
) -> T:
ensure_dir(output_dir)
profile = cProfile.Profile()
profile.enable()
try:
return fn(*args, **kwargs)
finally:
profile.disable()
summary = render_cprofile_summary(profile, sort_by=sort_by, limit=limit)
(output_dir / f"{label}.txt").write_text(summary)
def _stable_float(value: float | int | None) -> float | None:
if value is None:
return None
return round(float(value), 8)
def _tensor_signature(tensor: torch.Tensor) -> dict[str, Any]:
cpu_tensor = tensor.detach().cpu()
if cpu_tensor.numel() == 0:
stats = {"sum": None, "mean": None, "std": None, "min": None, "max": None}
else:
stats_tensor = (
cpu_tensor.to(torch.float64) if cpu_tensor.is_floating_point() else cpu_tensor.to(torch.int64)
)
stats = {
"sum": _stable_float(stats_tensor.sum().item()),
"mean": _stable_float(stats_tensor.float().mean().item()),
"std": _stable_float(stats_tensor.float().std(unbiased=False).item())
if cpu_tensor.numel() > 1
else 0.0,
"min": _stable_float(stats_tensor.min().item()),
"max": _stable_float(stats_tensor.max().item()),
}
hash_tensor = cpu_tensor.float() if cpu_tensor.dtype == torch.bfloat16 else cpu_tensor
digest = hashlib.sha256(hash_tensor.contiguous().numpy().tobytes()).hexdigest()
return {
"shape": list(cpu_tensor.shape),
"dtype": str(cpu_tensor.dtype),
"numel": cpu_tensor.numel(),
"sha256": digest,
**stats,
}
def _summarize_forward_value(value: Any) -> Any:
if isinstance(value, torch.Tensor):
return _tensor_signature(value)
if isinstance(value, dict):
return {key: _summarize_forward_value(val) for key, val in value.items()}
if isinstance(value, (list, tuple)):
return [_summarize_forward_value(item) for item in value]
if isinstance(value, (str, int, float, bool)) or value is None:
return value
return repr(value)
def _hash_payload(payload: Any) -> str:
return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
def _build_reference_batch(dataset: Any, batch_size: int) -> Any:
if len(dataset) == 0:
raise ValueError("Cannot build a reference batch from an empty dataset.")
indices = [idx % len(dataset) for idx in range(batch_size)]
samples = [dataset[idx] for idx in indices]
return default_collate(samples)
def write_deterministic_forward_artifacts(
*,
policy: Any,
dataset: Any,
batch_size: int,
preprocessor: Any,
output_dir: Path,
device_type: str,
) -> None:
reference_batch = preprocessor(_build_reference_batch(dataset, batch_size))
activities = [torch.profiler.ProfilerActivity.CPU]
if device_type == "cuda":
activities.append(torch.profiler.ProfilerActivity.CUDA)
was_training = policy.training
policy.eval()
with torch.random.fork_rng(devices=[] if device_type != "cuda" else None):
torch.manual_seed(0)
if device_type == "cuda":
torch.cuda.manual_seed_all(0)
with torch.no_grad(), torch.profiler.profile(activities=activities) as profiler:
loss, output_dict = policy.forward(reference_batch)
if was_training:
policy.train()
operator_entries = []
for event in profiler.key_averages():
entry = {
"key": event.key,
"count": event.count,
"cpu_time_total_us": _stable_float(getattr(event, "cpu_time_total", None)),
}
if device_type == "cuda":
entry["self_cuda_time_total_us"] = _stable_float(getattr(event, "self_cuda_time_total", None))
operator_entries.append(entry)
operator_entries = sorted(operator_entries, key=lambda item: item["key"])
output_summary = {
"loss": _summarize_forward_value(loss),
"output_dict": _summarize_forward_value(output_dict),
}
payload = {
"seed": 0,
"reference_batch_size": batch_size,
"operator_fingerprint": _hash_payload([(entry["key"], entry["count"]) for entry in operator_entries]),
"output_fingerprint": _hash_payload(output_summary),
"operators": operator_entries,
"outputs": output_summary,
}
(output_dir / "deterministic_forward.json").write_text(json.dumps(payload, indent=2, sort_keys=True))
table_sort = "self_cuda_time_total" if device_type == "cuda" else "cpu_time_total"
write_profiler_table(profiler, output_dir / "deterministic_forward_ops.txt", sort_by=table_sort)
def _summary(values: list[float]) -> dict[str, float] | dict[str, None]:
if not values:
return {"count": 0, "mean": None, "median": None, "min": None, "max": None}
return {
"count": len(values),
"mean": statistics.fmean(values),
"median": statistics.median(values),
"min": min(values),
"max": max(values),
}
@dataclass
class StepTimingCollector:
forward_s: list[float] = field(default_factory=list)
backward_s: list[float] = field(default_factory=list)
optimizer_s: list[float] = field(default_factory=list)
total_update_s: list[float] = field(default_factory=list)
dataloading_s: list[float] = field(default_factory=list)
memory_timeline: list[dict[str, float | int]] = field(default_factory=list)
def record(
self,
*,
forward_s: float,
backward_s: float,
optimizer_s: float,
total_update_s: float,
) -> None:
self.forward_s.append(forward_s)
self.backward_s.append(backward_s)
self.optimizer_s.append(optimizer_s)
self.total_update_s.append(total_update_s)
def record_dataloading(self, dataloading_s: float) -> None:
self.dataloading_s.append(dataloading_s)
def record_memory(self, *, step: int, allocated_bytes: int, reserved_bytes: int) -> None:
self.memory_timeline.append(
{
"step": step,
"allocated_bytes": allocated_bytes,
"reserved_bytes": reserved_bytes,
}
)
def to_dict(self) -> dict[str, Any]:
return {
"forward_s": _summary(self.forward_s),
"backward_s": _summary(self.backward_s),
"optimizer_s": _summary(self.optimizer_s),
"total_update_s": _summary(self.total_update_s),
"dataloading_s": _summary(self.dataloading_s),
"memory_timeline": self.memory_timeline,
}
def write_json(self, output_path: Path, extra: dict[str, Any] | None = None) -> None:
payload = self.to_dict()
if extra:
payload.update(extra)
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(json.dumps(payload, indent=2, sort_keys=True))