Fold CP

Plugin fold-cp from marketplace boltz-cp in NVIDIA-BioNeMo/boltz-cp. Commands read from the repo's manifest — 1 month ago.

Skills (10)

• benchmark

  Benchmark a CP inference or training workflow to find the maximum token count that fits at a given CP size and to record end-to-end walltime. Sweeps input size (tokens N, atoms, MSA depth S) over real or synthesized data, ramps N until OOM to report the max, times fwd (inference) or fwd+bwd+step (tr

• build_infra

  Probe and establish the distributed test infrastructure for CP development. Inventories local GPUs (count, model, memory, NVLink topology), checks the software stack (Python, PyTorch+CUDA, NCCL, torch.distributed), and runs shipped smoke tests for batch_isend_irecv P2P, all_gather, all_reduce, reduc

• cpize_model_workflow

  Orchestrate the END-TO-END integration of context parallelism into a custom co-folding / structure-prediction model: turn the whole effort into a prioritized, dependency-sorted worklist, then drive it phase by phase — map the model (learn_context), verify infra (build_infra), shard data features (sh

• dispatch_work

  Orchestrate a team of agents to integrate many CP modules and tests in parallel. Builds the work-list from the module map, derives a data-flow dependency graph, topologically sorts it into GPU-slot-bounded waves, and spawns a coder+reviewer pair per work item. Reviewers enforce serial-as-ground-trut

• dist_lifecycle

  Stand up the distributed-model lifecycle for a CP training/inference run: device placement before DTensor wrapping, the all-trainable-params-are-DTensors invariant (with placeholder/freeze for unimplemented modules), checkpoint save (DTensor → plain) and load (realign via the live state_dict templat

• dtensor_modules

  Implement DTensor-based context-parallel model modules that mirror a serial reference. Maps each serial layer/module to its Boltz-CP counterpart with exact input/output shapes, placements, collectives, and backward memory budget; writes the CP implementation following the autograd.Function conventio

• learn_context

  Explore a custom co-folding / structure-prediction model to understand its inference and training workflows, entry points, training framework (PyTorch Lightning vs DeepSpeed), data-feature format, and featurization pipeline — then map those components onto the existing Boltz-CP context-parallel refe

• mem_profile

  Memory-profile a context-parallel (CP) inference (or training) workflow with the PyTorch CUDA caching-allocator history, then attribute the top-N memory peaks to specific modules and lines of code. Wraps the end-to-end forward in torch.cuda.memory._record_memory_history() + _dump_snapshot() (the sam

• nsys_profile

  Profile a context-parallel (CP) inference (or training) workflow with NVIDIA Nsight Systems (nsys). Locates or installs the nsys binary, mocks random features + small Glorot-init weights (or loads a real checkpoint), wraps the end-to-end forward in a torchrun launcher, runs nsys with the right CLI o

• shard_data_feats

  Implement distributed data-feature sharding for CP: assign DTensor placements to every model feature, build the placement-definition dictionary, implement atom-feature pack/pad/scatter, per-shard and cross-axis divisibility padding, the DTensor Dataset / DataLoader / DataModule and collate, and the

Hooks (2)

• SessionStart

  cat "${CLAUDE_PLUGIN_ROOT}/hooks/RULES.md"

• PreToolUse

  python3 "${CLAUDE_PLUGIN_ROOT}/hooks/guard_serial_and_tests.py"

Similar plugins