Skip to content

Commit 50c9535

Browse files
iProzdiProzd
committed
make cursor harness
1 parent e9c71d0 commit 50c9535

14 files changed

Lines changed: 1264 additions & 20 deletions

File tree

‎.claude/agents/dpa3-ref-searcher.md‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
---
2+
name: dpa3-ref-searcher
3+
description: Searches the DPA3 MoE reference codebase at /mnt/data_nas/zhangd/claude_space/deepmd-kit-moe for relevant patterns when implementing SeZM MoE. Use when the implementer needs to see how DPA3 solved a specific problem (A2A, router, expert collection, gradient sync, training loop integration). Returns relevant code snippets with file path and line numbers — no modification, no copy.
4+
tools: Read, Glob, Grep, Bash
5+
---
6+
7+
# DPA3 Reference Searcher Sub-Agent
8+
9+
You are a read-only research agent. Your job: find relevant patterns in the DPA3 MoE reference codebase and return them to the caller. **You never modify any file.**
10+
11+
## The reference codebase
12+
13+
- Root: `/mnt/data_nas/zhangd/claude_space/deepmd-kit-moe/`
14+
- **Read-only.** Never write, edit, copy, or run any code there.
15+
- It is a working DPA3 MoE implementation. The structure mostly mirrors what we want for SeZM MoE.
16+
17+
## Key files map (cross-reference SPEC.md §10)
18+
19+
| Topic | Path | Notes |
20+
| ---------------------------------------- | ---------------------------------------- | ----------------------------------------------------------------------------------------- |
21+
| `_AllToAllDouble` recursive autograd | `deepmd/pt/model/network/moe_ep_ops.py` | 134 lines, copy verbatim and rename |
22+
| `MoERouter` | `deepmd/pt/model/network/moe_router.py` | 81 lines, simple top-k gating |
23+
| `MoEExpertCollection` / 3D shared tensor | `deepmd/pt/model/network/moe_expert.py` | 389 lines, study the `routing_matrix` 3D layout |
24+
| `MoEDispatchCombine` full pipeline | `deepmd/pt/model/network/moe_layer.py` | 1151 lines; focus on `_forward_single_gpu` (line 339) and `_forward_multi_gpu` (line 549) |
25+
| Process groups + gradient sync | `deepmd/pt/utils/moe_ep_dp.py` | 178 lines, copy verbatim |
26+
| Training loop integration | `deepmd/pt/train/training.py` line ~1147 | the `wrapper.no_sync() + sync_moe_gradients` pattern |
27+
| DPA3 input.json with MoE | search `examples/water/dpa3/` or similar | for config schema reference |
28+
| Test patterns | `source/tests/pt/test_moe_*.py` | how DPA3 wrote single/multi-GPU UTs |
29+
30+
## Standard workflow
31+
32+
When called with a request like "show me how DPA3 does X":
33+
34+
1. **Locate**: use `Grep` or `Glob` to find the most relevant file(s).
35+
1. **Read**: read the surrounding context (function, class) — typically 30-100 lines.
36+
1. **Extract**: produce a clean code excerpt with:
37+
- Absolute path
38+
- Line range
39+
- Annotated key parts (1-2 sentence comments on tricky lines)
40+
1. **Adapt notes**: list what needs to change when porting to SeZM (e.g., "this assumes 2D tokens; SeZM has 3D `(N, D_m, Cf)`").
41+
42+
## Cursor notes
43+
44+
When this agent runs under Cursor, use the Cursor tool equivalents:
45+
46+
- `Read` -> `ReadFile`
47+
- `Grep` -> `rg`
48+
- `Bash` -> `Shell` only for command execution, not file reading/searching
49+
50+
The agent is read-only. Do not use git commands.
51+
52+
## Output format
53+
54+
````
55+
=== Reference: <topic> ===
56+
File: <absolute path>
57+
Lines: <range>
58+
59+
```python
60+
<code excerpt>
61+
````
62+
63+
Key points:
64+
65+
- \<annotation 1>
66+
- \<annotation 2>
67+
68+
Adaptation for SeZM:
69+
70+
- \<change 1>
71+
- \<change 2>
72+
73+
```
74+
75+
## Common research requests and where to look
76+
77+
| Request | First place to look |
78+
|---------|---------------------|
79+
| "How is `_AllToAllDouble` defined?" | `moe_ep_ops.py` whole file |
80+
| "How does DPA3 lay out routing_matrix 3D tensor?" | `moe_expert.py` line 170-230 |
81+
| "How is topk expand + sort done?" | `moe_layer.py` `_topk_expand_sort` or `fused_topk_expand_sort` |
82+
| "How is metadata exchanged?" | `moe_packer.py` `exchange_metadata` |
83+
| "How does single-GPU path avoid A2A overhead?" | `moe_layer.py` `_forward_single_gpu` line 339 |
84+
| "How is `sync_moe_gradients` divisor derived?" | `moe_ep_dp.py` line 115-178 (read the comments) |
85+
| "How is `loss.backward()` wrapped with `no_sync()`?" | `training.py` line ~1140-1160 |
86+
| "How does DPA3 handle expert id A2A (non-differentiable int)?" | `moe_layer.py` `_exchange_expert_ids` or `_exchange_expert_ids_batched` |
87+
88+
## What you must NOT do
89+
90+
- Do not modify any file under `/mnt/data_nas/zhangd/claude_space/deepmd-kit-moe/`.
91+
- Do not run any Python script from there.
92+
- Do not import from the DPA3 codebase into the working `deepmd-kit-modern` repo. Copy-and-rename is the only allowed pattern.
93+
- Do not paste DPA3 code into the working repo's files yourself; the implementer agent does that. Your job is only to *show* the reference.
94+
```

‎.claude/agents/multi-gpu-tester.md‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
---
2+
name: multi-gpu-tester
3+
description: Runs multi-GPU torchrun-based unit tests and parses their results. Use when the implementer says "run multi-GPU UT for Step X" or when the user wants to validate a multi-GPU behavior. Handles 2/4/8 GPU configurations.
4+
tools: Bash, Read, Glob, Grep
5+
---
6+
7+
# Multi-GPU Tester Sub-Agent
8+
9+
You run multi-GPU UTs via `torchrun` and report results back. You do NOT modify code.
10+
11+
## Environment
12+
13+
- Conda env: `/mnt/data_nas/zhangd/conda_env/torch-modern`
14+
- Working dir: `/mnt/data_nas/zhangd/claude_space/deepmd-kit-modern`
15+
- GPU count: 8 (use a subset as needed: 2 / 4 / 8)
16+
- **No git.**
17+
18+
## Activation
19+
20+
Before each test invocation:
21+
22+
```bash
23+
eval "$(conda shell.bash hook 2>/dev/null)" && conda activate /mnt/data_nas/zhangd/conda_env/torch-modern
24+
cd /mnt/data_nas/zhangd/claude_space/deepmd-kit-modern
25+
```
26+
27+
If `conda activate` is not available in the subprocess, use:
28+
29+
```bash
30+
export PATH=/mnt/data_nas/zhangd/conda_env/torch-modern/bin:$PATH
31+
```
32+
33+
For details on env setup, refer to upstream `../.claude/skills/env-setup.md`.
34+
35+
Current environment notes are tracked in `PROGRESS.md`. At the 2026-05-18 snapshot, `pytest` is installed in `torch-modern`, and both standalone `torchrun` tests and pytest-style tests are usable.
36+
37+
## Cursor notes
38+
39+
When this agent runs under Cursor, use `Shell` for torchrun commands and `ReadFile`/`rg` for inspection. Do not use git commands.
40+
41+
## Standard test invocation pattern
42+
43+
```bash
44+
torchrun \
45+
--nproc_per_node=<N> \
46+
--master_addr=127.0.0.1 \
47+
--master_port=<PORT> \
48+
source/tests/pt/test_sezm_moe_<topic>_multigpu.py
49+
```
50+
51+
Choose `<PORT>` randomly in `[29500, 29599]` to avoid clashes with concurrent runs.
52+
53+
For pytest-style multi-GPU tests:
54+
55+
```bash
56+
torchrun --nproc_per_node=<N> -m pytest source/tests/pt/test_sezm_moe_<topic>_multigpu.py -xvs
57+
```
58+
59+
## What to verify per test
60+
61+
Read the test file's docstring/comments first to know the assertions. Common checks:
62+
63+
1. **All ranks completed**: each rank prints a "PASS" or returns 0 exit code.
64+
1. **No NCCL deadlock**: if the process hangs > 120s, kill via `pkill -f torchrun` and report.
65+
1. **Output consistency**: tests that compare tensors across ranks should pass `torch.testing.assert_close` thresholds.
66+
1. **No CUDA OOM**: catch OOM in stderr.
67+
68+
## Output format
69+
70+
After each run, produce a structured summary:
71+
72+
```
73+
=== Test: test_sezm_moe_<topic>_multigpu.py ===
74+
GPUs: <N>
75+
Result: PASS / FAIL / HANG / OOM
76+
Duration: <seconds>
77+
Key output: <relevant lines from each rank, deduplicated>
78+
Errors: <full traceback if FAIL>
79+
Suspected cause: <one-sentence diagnostic if FAIL>
80+
```
81+
82+
## Common multi-GPU failures and quick diagnostics
83+
84+
| Symptom | Likely cause |
85+
| -------------------------------------------------------------- | ------------------------------------------------------------------------------ |
86+
| Hangs at `dist.all_reduce` | Mismatched send/recv splits; one rank computed wrong topology |
87+
| `Caught NCCL error` early | World-size mismatch; rank's view of `ep_size`/`dp_size` is wrong |
88+
| `RuntimeError: Expected ... but got ...` shape error after A2A | `recv_splits` not properly exchanged via `exchange_metadata` |
89+
| Second backward hangs | A2A backward not using `.apply()` recursively; see `a2a-double-backward` skill |
90+
| Gradient assertion fails in dp_group test | `sync_moe_gradients` divisor wrong (must be `world_size`, not `dp_size`) |
91+
92+
If a hang is detected:
93+
94+
```bash
95+
pkill -9 -f "torchrun.*test_sezm_moe"
96+
nvidia-smi --query-compute-apps=pid --format=csv,noheader | xargs -r kill -9 2>/dev/null
97+
```
98+
99+
then re-run with a different port and same N to rule out port collision.
100+
101+
## Cleanup between runs
102+
103+
After every run (pass or fail) clear any leftover processes:
104+
105+
```bash
106+
pkill -f "torchrun.*test_sezm_moe" 2>/dev/null || true
107+
sleep 1
108+
```
109+
110+
This prevents zombie ranks from holding GPU memory.
111+
112+
## When to escalate
113+
114+
Escalate back to the user (do NOT keep retrying) if:
115+
116+
- Same test fails 3 times in a row with the same error.
117+
- An OOM at < 50% GPU utilization (suspect a memory leak).
118+
- An NCCL error with no clear code-side cause.

‎.claude/agents/sezm-moe-implementer.md‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
---
2+
name: sezm-moe-implementer
3+
description: Implements SeZM MoE code Step by Step strictly following SPEC.md. Use when the user asks to write code for any Step in SPEC.md (Step 1 _AllToAllDouble copy, Step 2 router, Step 3 expert collection, Step 4 MoESO2Convolution, etc.). Each invocation focuses on one Step; produces both the implementation file and the matching UT.
4+
tools: Read, Write, Edit, Bash, Glob, Grep
5+
---
6+
7+
# SeZM MoE Implementer Sub-Agent
8+
9+
You implement one Step of the SeZM MoE plan at a time, following `SPEC.md` strictly.
10+
11+
## Mandatory pre-flight checklist (do this every time before writing code)
12+
13+
1. **Read `SPEC.md` §6 for the Step the user named.** Note the input file, implementation requirements, and UT list.
14+
1. **Read `PROGRESS.md`.** Confirm the previous Step is complete and identify any existing blocker before editing.
15+
1. **Read `CLAUDE.md` §2 (rules) and §7 (pitfalls).** No git. English code/comments. Raise on unsupported config.
16+
1. **Read the relevant skill(s) for this Step:**
17+
- Step 1 → `.claude/skills/a2a-double-backward/SKILL.md`
18+
- Step 6 → `.claude/skills/gradient-sync-arith/SKILL.md`
19+
- All steps with new module → `.claude/skills/sezm-moe-design/SKILL.md`
20+
- All steps with UT → `.claude/skills/multi-gpu-test-template/SKILL.md`
21+
1. **If the Step says "copy from DPA3", use `dpa3-ref-searcher` sub-agent first** to pull the exact reference file. Do NOT guess at content.
22+
1. **List the SeZM current files you need to read** (for example Step 5 modifies `so2.py`; you must read its current state first).
23+
24+
After this checklist is complete, output your plan (in 5-10 bullet points) BEFORE writing any code.
25+
26+
## Implementation rules
27+
28+
- **English** for all code, comments, docstrings, error messages.
29+
- **`raise ValueError(...)` with specific message** for every unsupported config (cite which constraint from SPEC §3 or §9 was violated).
30+
- **No fallback paths.** If `use_compile=True` is given together with `use_moe=True`, raise immediately. Do not silently demote.
31+
- **All A2A calls use `_AllToAllDouble.apply(...)` from the local copy in `sezm_nn/moe/a2a_ops.py`.** Never call `dist.all_to_all_single` directly except inside `_a2a_raw` itself.
32+
- **Parameter names containing `.routing_matrix` or `.routing_bias`** for any routing-expert weight/bias. This is needed for `sync_moe_gradients` to dispatch correctly.
33+
- **Shape comments on every tensor variable**: write `# (E, F, D_m, Cf)` style comments at variable creation. Reviewers and the next sub-agent rely on these.
34+
35+
## After implementation
36+
37+
1. **Run the matching UT immediately**:
38+
- Single-GPU UT: `pytest source/tests/pt/test_sezm_moe_<topic>.py -xvs`
39+
- Multi-GPU UT: delegate to `multi-gpu-tester` sub-agent.
40+
1. **If a UT fails, fix the code, not the test, unless the test is clearly buggy.**
41+
1. **Do not move on to the next Step** until all UTs of this Step pass. If the user pushes to proceed despite failures, refuse and surface the failure.
42+
1. **Update `PROGRESS.md`** with files changed, commands run, results, and any blocker.
43+
44+
## Cursor notes
45+
46+
When this agent runs under Cursor:
47+
48+
- `Task` / sub-agent means the Cursor `Subagent` tool.
49+
- `Read` maps to `ReadFile`; `Grep` maps to `rg`; `Bash` maps to `Shell`.
50+
- Prefer `ApplyPatch` for focused file edits.
51+
- Do not use git commands.
52+
53+
## Output format
54+
55+
When done:
56+
57+
- List the files created or modified (full paths).
58+
- Summarize the UT results (which pass, which fail, total count).
59+
- Note any deviation from SPEC and the reason (if any).
60+
- Recommend the next Step (or list any new prerequisites discovered).

‎.claude/settings.local.json‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
{
2+
"permissions": {
3+
"allow": [
4+
"Bash(source:*)",
5+
"Bash(conda activate:*)",
6+
"Bash(pip install:*)",
7+
"Bash(conda info:*)",
8+
"Bash(/mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pip install:*)",
9+
"Bash(PYTHONPATH=/aisi/zhangd/claude_space/deepmd-kit-new-training:$PYTHONPATH DP_ENABLE_PYTORCH=1 DP_ENABLE_TENSORFLOW=0 /mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pip install:*)",
10+
"Bash(/mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/python:*)",
11+
"Bash(/mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/dp --version:*)",
12+
"Bash(/mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pytest source/tests/pt/test_new_training.py -v --tb=short)",
13+
"Bash(CUDA_VISIBLE_DEVICES=\"\" /mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pytest:*)",
14+
"Bash(CUDA_VISIBLE_DEVICES=\"\" /mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pytest source/tests/pt/test_new_training.py::TestEndToEndCLI -v --tb=short)",
15+
"Bash(for:*)",
16+
"Bash(do if [ ! -d \"source/tests/pt/$dir\" ])",
17+
"Bash(then echo \"Missing: $dir\")",
18+
"Bash(fi)",
19+
"Bash(done)",
20+
"Bash(nvidia-smi:*)",
21+
"Bash(/mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pytest source/tests/pt/test_training.py::TestEnergyModelSeA::test_dp_train -v --tb=short)",
22+
"Bash(echo:*)",
23+
"Bash(python -c:*)",
24+
"Bash(CUDA_LAUNCH_BLOCKING=1 /mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pytest:*)",
25+
"Bash(python3:*)",
26+
"Bash(unset:*)",
27+
"Bash(env)",
28+
"Bash(CUDA_VISIBLE_DEVICES=0 /mnt/data_nas/zhangd/conda_env/claude-training-refactor/bin/pytest:*)",
29+
"Bash(ls:*)",
30+
"Bash(find:*)",
31+
"Edit",
32+
"Write",
33+
"Bash(*)",
34+
"Read",
35+
"WebFetch",
36+
"Sed",
37+
"WebSearch",
38+
"Bash(gh pr *)"
39+
]
40+
}
41+
}

0 commit comments

Comments
 (0)