Skip to content

cli.md

Latest commit

 

History

History
143 lines (105 loc) · 9.17 KB

cli.md

File metadata and controls

143 lines (105 loc) · 9.17 KB
Raw

Command-line interface

Hydra composes the command-line configuration, and Pydantic validates it. The schemas are declared in config.py and registered in config_store.py.

Commands

The console scripts are gpuastar and gpuastar-aoti-export. The top-level command option selects the execution module:

  • astar_search (default): runs batch weighted A* search on a set of input states.
  • aoti_export: compiles a heuristic model checkpoint into an AOTInductor shared library package. The gpuastar-aoti-export console script runs the same exporter.
  • cpp_workflows: runs C++ build and solver workflows.

The defaults listed below are the values defined in the base gpuastar package. When the gpuastar[experiments] extra is installed, the overlays in command_defaults/ replace some of them with Rubik's Cube example values, which are noted where they apply. Explicit command line overrides always win.

Configuration groups

Options are grouped into namespaces and overridden with dot notation, for example algorithm.batch_size=5000.

1. Input settings (input)

Paths and index ranges for the search run. The base package has no environment, dataset, model, or results defaults, so the first four options must be set explicitly unless the gpuastar[experiments] extra provides them.

  • input.states: path to the serialized input states file. No base default. The experiments overlay sets experiments/src/gpuastar_experiments/data/cube3/test/data_0.pkl.
  • input.model: path to the neural network checkpoint. No base default. The experiments overlay sets experiments/src/gpuastar_experiments/saved_models/cube3/current/model_state_dict.pt.
  • input.env: registered name of the search environment. No base default. The experiments overlay sets cube3.
  • input.results_dir: output directory for search results. No base default. The experiments overlay sets results/cube3.
  • input.start_idx: first state index to read from the input file. Default is 0.
  • input.max_states: maximum number of states to process, -1 processes all states. Default is -1.

2. Algorithm settings (algorithm)

Search parameters for the weighted evaluation.

  • algorithm.batch_size: number of states popped from the open list and expanded per iteration. Base default is 1. The experiments overlay sets 10000.
  • algorithm.weight: multiplier $w$ applied to the heuristic value $h(s)$ in $f(s) = g(s) + w \cdot h(s)$. Base default is 1.0. The experiments overlay sets 0.6.
  • algorithm.instance_batch_size: number of independent search instances run in parallel. Default is 1.
  • algorithm.hash_capacity_factor: capacity of the closed-set hash table relative to the expected number of states. Default is 4.0.
  • algorithm.capacity_mode: resizing policy for the closed-set hash table, adaptive or fixed. Default is adaptive.
  • algorithm.hashtable_algorithm: closed-set implementation, python_dict or linear_probing. Default is python_dict.
  • algorithm.solution_move_source: solution path reconstruction, infer retraces parent pointers and stored reads stored moves. Default is infer.
  • algorithm.duplicate_update_mode: handling of rediscovered states, in_place updates the existing slot and new_slot inserts a new entry. Default is in_place.
  • algorithm.goal_stop_mode: termination criterion, stop_before_goal halts when a goal is popped and expand_popped_batch finishes the current popped batch first. Default is stop_before_goal.
  • algorithm.expand_push_mode: child-materialization used by Python CPU A* (backend.language=python, backend.gpu_only=false) only. single_pass builds child nodes in one index-tracking pass used by the paper's reference implementation. flatten uses heavier flatten/unflatten reshaping that produces identical child nodes with extra per-batch work. Default is single_pass. Other backends accept only the default value.
  • algorithm.nnet_batch_size: batch size for heuristic model inference, 0 matches the expansion batch size. Base default is 0. The experiments overlay sets 10000.
  • algorithm.search_budget: limit on node expansions, 0 disables the limit. Default is 0.

3. Backend settings (backend)

Solver implementation and device selection.

  • backend.language: solver backend, one of python, cpp, cu, or auto. Default is python.
  • backend.gpu_only: if true, the entire search loop runs in device-resident tensors instead of host data structures. Default is false.
  • backend.device: execution device, for example cpu, cuda:0, or auto. Default is auto.
  • backend.heuristic_device: separate device for model inference. Falls back to backend.device when unset. Default is unset.

4. PyTorch settings (torch)

Compilation, precision, and distributed execution controls.

  • torch.compile_all: if true, enables torch.compile for all compile targets. Default is false.
  • torch.compile: torch.compile options for heuristic evaluation (fields enabled, backend, fullgraph, dynamic, options).
  • torch.compile_env: torch.compile options for the environment transition functions.
  • torch.compile_astar_closed: torch.compile options for closed-set matching.
  • torch.compile_mode: compilation mode, such as default, reduce-overhead, or max-autotune. Default is default.
  • torch.fsdp: enables Fully Sharded Data Parallel (FSDP2) parameter sharding. Default is false.
  • torch.fsdp_reshard_after_forward: whether sharded parameters are freed after each forward pass. Default is unset, which uses the FSDP2 default.
  • torch.tensor_parallel: enables tensor parallel sharding across GPUs. Default is false.
  • torch.dtensor: wraps model parameters as DTensors. Default is false.
  • torch.dtensor_placement: DTensor placement, replicate or shard_dim0. Default is replicate.
  • torch.distributed_sort_policy: state grouping before scattering to ranks, none, local, or global. Default is none.
  • torch.cuda_graphs: captures the static portion of the search loop as a CUDA graph. Default is false.
  • torch.cudagraph_warmup: warmup iterations executed before graph capture. Default is 2.
  • torch.amp: enables automatic mixed precision for inference. Default is false.
  • torch.amp_dtype: mixed precision dtype, float16 or bfloat16. Default is float16.
  • torch.tf32: enables TensorFloat-32 matrix math on NVIDIA Ampere and newer GPUs. Default is false.

5. Memory dtype settings (memory_dtype)

Data types of the tensor-backed search storage. Narrower types reduce memory traffic in device-resident modes.

  • memory_dtype.node_state: dtype of stored states. Default is int64.
  • memory_dtype.open_index: dtype of open-list indices. Default is int64.
  • memory_dtype.open_cost: dtype of open-list priorities. Default is float32.
  • memory_dtype.closed_key: dtype of closed-set keys. Default is int64.
  • memory_dtype.closed_cost: dtype of closed-set costs. Default is float32.
  • memory_dtype.parent_index: dtype of parent pointers. Default is int64.
  • memory_dtype.parent_move: dtype of stored parent moves. Default is int64.
  • memory_dtype.path_cost: dtype of the path cost $g(s)$. Default is float32.
  • memory_dtype.heuristic: dtype of the heuristic estimate $h(s)$. Default is float32.
  • memory_dtype.total_cost: dtype of the total cost $f(s)$. Default is float32.

6. Profile settings (profile)

Profiler and timing instrumentation.

  • profile.enabled: activates the PyTorch profiler. Default is false.
  • profile.operation_breakdown: records timing for the individual search operations (pop, expand, heuristic, push, closed-set check). Default is false.
  • profile.cuda_sync_timing: synchronizes CUDA streams around measured intervals so that durations include kernel execution. Default is false.
  • profile.cpp_trace and profile.cpp_summary: enable the timers built into the C++ solvers. Defaults are false.

7. Metrics settings (metrics)

  • metrics.json_path: destination of the metrics file. Default is <results_dir>/search_metrics.json.
  • metrics.no_json: disables writing the metrics file. Default is false.

8. Verbosity (verbose)

  • verbose: console output level from 0 (quiet) upward. Default is 1.

Examples

Search with the Python backend on GPU

gpuastar command=astar_search \
  backend.language=python \
  backend.device=cuda:0 \
  algorithm.batch_size=1000 \
  algorithm.weight=0.8

C++ LibTorch GPU solver with explicit cost dtypes

gpuastar command=astar_search \
  backend.language=cpp \
  backend.gpu_only=true \
  memory_dtype.open_cost=float32 \
  memory_dtype.closed_cost=float32

Exporting a model to an AOTInductor package

gpuastar-aoti-export \
  export.model_dir=experiments/src/gpuastar_experiments/saved_models/cube3/current \
  export.env=cube3 \
  export.max_dynamic_batch=10000