Diff Coverage

    group_size = dist.get_world_size(group)
    gather = _functional_all_gather if torch.is_grad_enabled() else _functional_all_gather_forward
    gathered = gather(data, 0, group)
    if hasattr(gathered, 'wait'):
        gathered = gathered.wait()
    if concat_dim == 0:
        return gathered
    # Reshape from [group_size * d0, d1, ...] to [d0 * group_size, d1, ...]
    # by moving the group_size factor from dim 0 to concat_dim.

        raise ValueError("Compiled differentiable reduce-scatter supports only sum or string 'avg'")
    scatter = _functional_reduce_scatter if torch.is_grad_enabled() else _functional_reduce_scatter_forward
    result = scatter(data, "sum", axis, group)
    if hasattr(result, 'wait'):
        result = result.wait()
    # 'avg' maps to SUM in _OP_MAP, so the division stays manual.
    if isinstance(op, str) and op == 'avg':
        result = result / dev_num
    return result

    config["num_hidden_layers"] = layers
    if "max_window_layers" in config:
        config["max_window_layers"] = min(config["max_window_layers"], layers)
    if isinstance(config.get("layer_types"), list):
        config["layer_types"] = config["layer_types"][:layers]
    if isinstance(config.get("mlp_only_layers"), list):
        config["mlp_only_layers"] = [index for index in config["mlp_only_layers"] if index < layers]
    (output / "config.json").write_text(json.dumps(config, indent=2), encoding="utf-8")
    for path in source.iterdir():

    (output / "config.json").write_text(json.dumps(config, indent=2), encoding="utf-8")
    for path in source.iterdir():
        if path.is_file() and path.name.startswith(("tokenizer", "vocab", "merges", "generation_config",
                                                   "special_tokens", "chat_template")):
            shutil.copy2(path, output / path.name)
    manifest = {"source": str(source.resolve()), "layers": layers,
                "source_config_sha256": _file_digest(source / "config.json"),
                "source_index_sha256": _file_digest(source / "model.safetensors.index.json"), "shards": []}
    exported = {"metadata": {"total_size": 0}, "weight_map": {}}

             (output / "config.json", "output_config_sha256"),
             (output / "model.safetensors.index.json", "output_index_sha256"))
    for path, key in paths:
        if _file_digest(path) != manifest[key]:
            raise ValueError(f"Checkpoint metadata differs from crop manifest: {path}")
    for shard in manifest["shards"]:
        if _file_digest(output / shard["output_file"]) != shard["sha256"]:
            raise ValueError(f"Cropped checkpoint shard checksum failed: {shard['output_file']}")

def main() -> None:
    """Prepare a separate artifact without modifying the pretrained source."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--layers", type=int, default=4)
    parser.add_argument("--reuse", action="store_true", help="Verify and reuse an existing crop")
    args = parser.parse_args()
    if args.reuse and args.output.exists():
        verify_checkpoint(args.source, args.output, args.layers)
    else:
        prepare_checkpoint(args.source, args.output, args.layers)


if __name__ == "__main__":
    main()

def run(config: dict[str, Any], report_path: Path) -> None:
    """Observe production diagnostics without changing tokens, rewards or numerical paths."""
    if config["train"]["max_steps"] != 2 or config["logging"]["log_steps"] != 1:
        raise ValueError("This example requires max_steps=2 and log_steps=1")
    report_path.parent.mkdir(parents=True, exist_ok=True)
    if int(os.environ.get("RANK", "0")) == 0 and report_path.exists():
        report_path.unlink()
    if int(os.environ.get("RANK", "0")) == 0:

    # This example adds an observation gate at the existing pre-update boundary.
    trainer.targets.prepare = _checked_prepare
    trainer.train()
    if trainer.state.global_step != 2 or [check["step"] for check in checks] != [1, 2]:
        raise RuntimeError("The normal RL loop did not complete and verify both updates")
    if int(os.environ.get("RANK", "0")) == 0:
        report = {"passed": True, "completed_steps": 2, "checks": checks,
                  "training": config["train"]["accelerator"],
                  "rollout": config["rollout"]["vllm"],

def main() -> None:
    """Load the explicit recipe and launch the ordinary synchronous trainer."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("config")
    parser.add_argument("--report", type=Path, default=Path("/results/run/consistency.json"))
    args, overrides = parser.parse_known_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
    run(load_config(args.config, overrides), args.report)


if __name__ == "__main__":
    main()