From ca1d829902d31211a6552a573649aa8656aa09f8 Mon Sep 17 00:00:00 2001 From: lxd-cumt <1141051934@qq.com> Date: Fri, 15 May 2026 14:05:43 +0800 Subject: [PATCH] Add upgrade skills for TransformerEngine-FL upstream sync --- .claude/skills | 1 + skills/e2e-stage-manager/SKILL.md | 350 +++++++ skills/te-fl-upstream-sync/SKILL.md | 130 +++ .../phases/01-setup-and-analysis.md | 184 ++++ .../phases/02-merge-and-integrate.md | 962 ++++++++++++++++++ .../phases/03-patch-and-verify.md | 432 ++++++++ .../phases/04-test-and-finalize.md | 443 ++++++++ .../references/cicd-pipeline.md | 258 +++++ .../scripts/generate_sync_report.sh | 116 +++ .../scripts/validate_plugin.sh | 142 +++ 10 files changed, 3018 insertions(+) create mode 120000 .claude/skills create mode 100644 skills/e2e-stage-manager/SKILL.md create mode 100644 skills/te-fl-upstream-sync/SKILL.md create mode 100644 skills/te-fl-upstream-sync/phases/01-setup-and-analysis.md create mode 100644 skills/te-fl-upstream-sync/phases/02-merge-and-integrate.md create mode 100644 skills/te-fl-upstream-sync/phases/03-patch-and-verify.md create mode 100644 skills/te-fl-upstream-sync/phases/04-test-and-finalize.md create mode 100644 skills/te-fl-upstream-sync/references/cicd-pipeline.md create mode 100644 skills/te-fl-upstream-sync/scripts/generate_sync_report.sh create mode 100644 skills/te-fl-upstream-sync/scripts/validate_plugin.sh diff --git a/.claude/skills b/.claude/skills new file mode 120000 index 0000000000..42c5394a18 --- /dev/null +++ b/.claude/skills @@ -0,0 +1 @@ +../skills \ No newline at end of file diff --git a/skills/e2e-stage-manager/SKILL.md b/skills/e2e-stage-manager/SKILL.md new file mode 100644 index 0000000000..0cb7f1271c --- /dev/null +++ b/skills/e2e-stage-manager/SKILL.md @@ -0,0 +1,350 @@ +--- +name: e2e-stage-manager +description: Manage end-to-end training test stages for FlagScale/TransformerEngine experiments. Use this skill when the user wants to create a unified stage configuration, migrate between stages (stage9/10/11), generate new stage configs, run batch training tests, or clean up old stage directories. Trigger on phrases like "create unified stage", "migrate stage configs", "run e2e tests", "clean up stages", "generate stage configs", or "consolidate training runs". +--- + +# E2E Stage Manager + +Manage end-to-end training test stages for FlagScale/TransformerEngine experiments. This skill helps you create unified stage configurations, migrate between stages, run batch tests, and clean up old runs. + +## Context + +The user's repo contains training test stages (stage9_runs, stage10_runs, stage11_runs) that test different TransformerEngine implementations: +- **Models**: Qwen3-32b, DeepSeek-V3 +- **TE implementations**: flagos, reference, vendor +- **Attention backends**: flash, fused, unfused (12 configs total per stage) + +Stage configurations are stored in `FlagScale/run_configs/stage{N}_*/` as Hydra YAML configs. + +## Stage Evolution + +**Stage 9 (Legacy):** +- Contains `legacy_tokenizer: true` in tokenizer config +- Missing `eval_interval` in model config +- `te_fl_prefer` parameter positioned after optimizer section + +**Stage 10+ (Current Standard):** +- Added `eval_interval: 1000` to model config +- Removed `legacy_tokenizer: true` from tokenizer config +- Moved `te_fl_prefer` parameter position (after `transformer_impl`) +- Re-runs use incremented stage numbers (stage11, stage12, ...) with identical format + +The latest deployed configs are stage11. Use Stage 10 format for all new stages. + +## Core Operations + +### 1. Create Unified Stage + +Generate a new unified stage configuration based on stage10 standard. + +**Steps:** +1. Create new stage directory structure: `FlagScale/run_configs/stage{N}_*/` +2. For each configuration (model × implementation × backend): + - Copy stage10 YAML structure + - Update `exp_name` to new stage number + - Update `exp_dir` path + - Ensure `eval_interval: 1000` is present + - Ensure `legacy_tokenizer` is removed +3. Create run directory: `stage{N}_runs/` + +**Config naming pattern:** +``` +stage{N}_te_fl_prefer-{impl}__attention_backend-{backend}/train.yaml +stage{N}_te_fl_prefer-{impl}__attention_backend-{backend}__deepseek_v3/train.yaml +``` + +Where: +- `{impl}` = flagos | reference | vendor +- `{backend}` = flash | fused | unfused + +### 2. Run Batch Tests + +Launch all 12 training configurations for a stage. + +**Steps:** +1. Read the stage's config directory +2. For each config YAML: + - Launch via FlagScale runner: `cd FlagScale && python -m flagscale.launcher.runner ` + - Capture output to `stage{N}_runs/logs/{model}_{impl}_{backend}.log` + - Track exit codes +3. Generate summary: `stage{N}_runs/logs/summary.txt` with PASS/FAIL for each config + +**Parallel execution:** +Run configs in parallel when possible (different models/configs don't conflict). + +### 3. Clean Up Old Stages + +Remove old stage directories while preserving important artifacts. + +**Steps:** +1. Ask user which stages to remove (default: keep only latest) +2. For each stage to remove: + - Archive logs to `backup/stage{N}_logs.tar.gz` + - Remove `stage{N}_runs/` directory + - Remove `FlagScale/run_configs/stage{N}_*/` configs +3. Report disk space freed + +**Safety:** Always confirm before deletion. Preserve summary.txt and final logs. + +### 4. Compare Stages + +Compare configurations or results between two stages. + +**Config diff:** +```bash +diff -r FlagScale/run_configs/stage9_* FlagScale/run_configs/stage10_* +``` + +**Results comparison:** +- Parse summary.txt from both stages +- Show pass/fail differences +- Compare timing/memory metrics if available + +### 5. Migrate Stage Configs + +Upgrade old stage configs to new standard (e.g., stage9 → stage10 format). + +**Steps:** +1. For each config in old stage: + - Load YAML + - Add `eval_interval: 1000` under `model:` + - Remove `legacy_tokenizer: true` from `data.tokenizer:` + - Reorder `te_fl_prefer` to come after `transformer_impl` + - Update stage number in paths +2. Write to new stage config directory +3. Validate YAML syntax + +## File Structure + +``` +repo_update/ +├── FlagScale/ +│ └── run_configs/ +│ ├── stage9_te_fl_prefer-flagos__attention_backend-fused/ +│ │ ├── train.yaml # Main config +│ │ └── train/32b.yaml # Model-specific overrides +│ └── stage10_te_fl_prefer-flagos__attention_backend-fused/ +│ ├── train.yaml +│ └── train/32b.yaml +├── stage9_runs/ +│ ├── Qwen3-32b-Stage9-flagos-fused/ +│ │ ├── checkpoints/ +│ │ ├── logs/ +│ │ ├── tensorboard/ +│ │ └── wandb/ +│ └── logs/ +│ └── summary.txt +└── stage10_runs/ + └── logs/ + └── summary.txt +``` + +## Config Template (Stage 10 Standard) + +### Qwen3-32b Config + +```yaml +defaults: + - _self_ + - train: 32b + +experiment: + exp_name: Qwen3-32b-Stage{N}-{impl}-{backend} + seed: 42 + save_steps: 10000 + load: null + exp_dir: /stage{N}_runs/${experiment.exp_name} + ckpt_format: torch + task: + type: train + backend: megatron + entrypoint: flagscale/train/megatron/train_gpt.py + runner: + per_node_task: false + no_shared_fs: false + rdzv_backend: static + hostfile: /host_single + ssh_port: 7878 + cmds: + before_start: ulimit -n 1048576 && source + envs: + LOGLEVEL: "INFO" + CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7" + CUDA_DEVICE_MAX_CONNECTIONS: 1 + +action: run + +hydra: + run: + dir: ${experiment.exp_dir}/hydra +``` + +### Model Config (train/32b.yaml) + +```yaml +system: + no_shared_fs: ${experiment.runner.no_shared_fs} + num_workers: 2 + tensor_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + precision: + bf16: true + attention_softmax_in_fp32: true + accumulate_allreduce_grads_in_fp32: true + logging: + log_interval: 1 + tensorboard_log_interval: 1 + wandb_project: ${experiment.exp_name} + wandb_exp_name: ${experiment.exp_name} + log_timers_to_tensorboard: true + log_validation_ppl_to_tensorboard: true + log_throughput: true + log_params_norm: true + log_num_zeros_in_grad: true + log_memory_to_tensorboard: true + checkpoint: + save_interval: ${experiment.save_steps} + load: ${experiment.load} + ckpt_format: ${experiment.ckpt_format} + +model: + transformer_impl: transformer_engine + te_fl_prefer: {impl} + attention_backend: {backend} + num_layers: 16 + hidden_size: 5120 + ffn_hidden_size: 25600 + num_attention_heads: 64 + kv_channels: 128 + group_query_attention: true + num_query_groups: 8 + seq_length: 4096 + max_position_embeddings: 40960 + norm_epsilon: 1.0e-06 + use_rotary_position_embeddings: true + rotary_base: 1000000 + swiglu: true + normalization: RMSNorm + qk_layernorm: true + init_method_std: 0.02 + attention_dropout: 0.0 + hidden_dropout: 0.0 + untie_embeddings_and_output_weights: true + no_position_embedding: true + no_rope_fusion: true + disable_bias_linear: true + seed: ${experiment.seed} + finetune: false + micro_batch_size: 1 + global_batch_size: 8 + eval_iters: 0 + eval_interval: 1000 # Stage 10 addition + train_iters: 20 + optimizer: + clip_grad: 1.0 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + lr_scheduler: + lr: 0.003 + min_lr: 0.0003 + lr_warmup_fraction: 0.1 + lr_decay_style: WSD + lr_wsd_decay_style: cosine + lr_wsd_decay_iters: 10 + +data: + reset_position_ids: true + reset_attention_mask: true + data_path: + split: 1 + no_mmap_bin_files: true + tokenizer: + tokenizer_type: QwenTokenizerFS # No legacy_tokenizer in Stage 10 + tokenizer_path: + vocab_size: 151851 + make_vocab_size_divisible_by: 64 +``` + +## Running Tests + +Launch a single config: +```bash +cd FlagScale +python -m flagscale.launcher.runner \ + --config-path ../run_configs/stage{N}_te_fl_prefer-{impl}__attention_backend-{backend} \ + --config-name train +``` + +Launch all configs for a stage: +```bash +# Create a runner script +for config in FlagScale/run_configs/stage{N}_*/train.yaml; do + config_dir=$(dirname $config) + config_name=$(basename $config_dir) + echo "=== [$config_name] START $(date) ===" >> stage{N}_runs/logs/${config_name}.log + cd FlagScale && python -m flagscale.launcher.runner \ + --config-path ../$config_dir \ + --config-name train \ + >> ../stage{N}_runs/logs/${config_name}.log 2>&1 + echo "=== [$config_name] EXIT_CODE=$? $(date) ===" >> stage{N}_runs/logs/${config_name}.log + cd .. +done +``` + +## Best Practices + +1. **Always use stage10 format for new stages** - includes eval_interval and removes legacy_tokenizer +2. **Test one config before batch runs** - verify environment and paths work +3. **Archive before cleanup** - preserve logs and summaries +4. **Use descriptive stage numbers** - increment sequentially (stage11, stage12, etc.) +5. **Check disk space** - each stage run can be 10-50GB depending on checkpoints + +## Common Issues + +**Issue**: Config not found +- Check `FlagScale/run_configs/` path exists +- Verify YAML syntax with `python -c "import yaml; yaml.safe_load(open('config.yaml'))"` + +**Issue**: CUDA OOM during parallel runs +- Run configs sequentially instead of parallel +- Reduce `micro_batch_size` or `global_batch_size` + +**Issue**: Checkpoint directory conflicts +- Ensure each config has unique `exp_dir` +- Clean old checkpoints before rerunning + +## Output Format + +When creating a unified stage, report: +``` +Created Stage {N} with 12 configurations: +✓ Qwen3-32b: flagos × [flash, fused, unfused] +✓ Qwen3-32b: reference × [flash, fused, unfused] +✓ Qwen3-32b: vendor × [flash, fused, unfused] +✓ DeepSeek-V3: [flagos, reference, vendor] × unfused + +Configs: FlagScale/run_configs/stage{N}_*/ +Run dir: stage{N}_runs/ + +To launch all tests: + bash scripts/run_stage{N}.sh +``` + +When running batch tests, show progress and final summary: +``` +Running Stage {N} tests... +[1/12] qwen3_flagos_flash... PASS (23.4s) +[2/12] qwen3_flagos_fused... PASS (24.1s) +... +[12/12] deepseek_v3_vendor_unfused... PASS (45.2s) + +Summary: 12/12 PASS (0 FAIL) +Results: stage{N}_runs/logs/summary.txt +``` diff --git a/skills/te-fl-upstream-sync/SKILL.md b/skills/te-fl-upstream-sync/SKILL.md new file mode 100644 index 0000000000..edc82bee99 --- /dev/null +++ b/skills/te-fl-upstream-sync/SKILL.md @@ -0,0 +1,130 @@ +--- +name: te-fl-upstream-sync +description: > + Manages the upstream sync workflow for the TransformerEngine-FL fork (flagos-ai/TransformerEngine-FL + forked from Nvidia/TransformerEngine). Handles creating dev branches aligned with upstream releases, + merging into main with conflict resolution, validating the custom plugin system (OP API interfaces, + CUDA patches, Python bindings), and running multi-level CI/CD verification. Use this skill whenever + the user mentions syncing upstream, updating from Nvidia/TransformerEngine, merging upstream releases + (e.g. release_v2.14), fork sync, repo update, pulling upstream changes, plugin validation, checking + plugin API, verifying CUDA patches, or any upstream integration workflow for TransformerEngine-FL. + Also trigger when the user references conflict resolution for plugin files, build system merges + involving setup.py/CMakeLists.txt with plugin targets, or running CI/CD pipelines after an upstream + merge. +--- + +# TransformerEngine-FL Upstream Sync Workflow + +You are guiding a developer through syncing their fork `flagos-ai/TransformerEngine-FL` with upstream +`Nvidia/TransformerEngine`. The fork's core value is a custom plugin system — every decision you make +must protect plugin functionality above all else. + +## Fork Architecture + +The fork adds these components on top of upstream: + +- **Plugin OP API interfaces**: `transformer_engine/plugin/` - Plugin systems +- **CUDA patches**: `transformer_engine/__init__.py` — Patches to upstream torch.cuda apis +- **Build integration**: `setup.py`, `CMakeLists.txt`, `pyproject.toml` — plugin compilation targets + woven into the upstream build system. +- **Other github workflow related** + +## Repo Detection Preamble + +Every stage below assumes you are inside the `TransformerEngine-FL` directory. Run this detection +snippet before any stage if the working directory is uncertain: + +```bash +if [ -d "TransformerEngine-FL" ]; then + cd TransformerEngine-FL +elif [ "$(basename $(pwd))" = "TransformerEngine-FL" ]; then + echo "Already in TransformerEngine-FL" +elif [ -d "../TransformerEngine-FL" ]; then + cd ../TransformerEngine-FL +else + echo "ERROR: TransformerEngine-FL directory not found in current or parent directory" + echo "Please clone the repo first, or cd to the correct location." + exit 1 +fi +echo "Working directory: $(pwd)" +``` + +--- + +## The Multi-Stage Sync Workflow + +This workflow is sequential. Each stage has a command the user can invoke, but you should also guide +them through the full flow when they ask to "sync upstream" or similar. + +### Phase Index + +| Phase | File | Stages | Description | +|-------|------|--------|-------------| +| 1 | [phases/01-setup-and-analysis.md](phases/01-setup-and-analysis.md) | Stage 1-2 | Repo Setup & Branch Preparation + Identify Plugin Changes | +| 2 | [phases/02-merge-and-integrate.md](phases/02-merge-and-integrate.md) | Stage 3-4 | Merge & Conflict Resolution + Plugin API Sync | +| 3 | [phases/03-patch-and-verify.md](phases/03-patch-and-verify.md) | Stage 5-7 | CUDA Patching + Stale Refs + Build Verify | +| 4 | [phases/04-test-and-finalize.md](phases/04-test-and-finalize.md) | Stage 8-10 | Tests + Merge to Main + FlagScale Training | + +Read the relevant phase file when the user enters that stage. Each phase is self-contained with +full instructions for its stages. + +--- + +## Sync Report (`/generate-sync-report`) + +After the workflow completes (success or failure), generate a report. Use the script at +`scripts/generate_sync_report.sh` to collect the data, or assemble manually: + +```markdown +# TransformerEngine-FL Upstream Sync Report + +## Summary +- **Date**: +- **Upstream Release**: release_v2.14 +- **Upstream Commit SHA**: +- **Status**: ✅ Complete / ❌ Failed at Stage N + +## Stage Results +| Stage | Status | Notes | +|-------|--------|-------| +| 1. Repo setup & branch preparation | ✅/❌ | | +| 2. Identify plugin changes | ✅/❌ | | +| 3. Merge & conflict resolution | ✅/❌ | N conflicts, P0: N, P1: N, P2: N | +| 4. Plugin API sync | ✅/❌ | N APIs added/modified/removed, N omissions found/fixed | +| 5. Patch CUDA hardcoding | ✅/❌ | | +| 6. Detect & fix stale references | ✅/❌ | N stale refs found/fixed | +| 7. Build & import verification | ✅/❌ | | +| 8. Unit & integration tests | ✅/❌ | CI script validation: N missing refs fixed, N tests added | +| 9. Merge to main | ✅/❌ | tree replacement merge, PR opened | +| 10. FlagScale training validation | ✅/❌ | N/M combinations passed (see batch comparison table) | + +## Conflicts Resolved + + +## Plugin System Status +- Plugin directory: ✅ intact +- OP API signatures: ✅ unchanged +- CUDA patches: ✅ present and applicable +- Build targets: ✅ present +- Python bindings: ✅ functional + +## CI/CD Results + + +## Rollback Info +- Merge commit: +- Rollback command: `git revert -m 1 ` +``` + +--- + +## Critical Rules + +These are non-negotiable because the fork's entire value proposition is the plugin system: + +1. **Plugin files are sacred.** Carefully auto-resolve a P0 conflict toward upstream. Always keep the + fork version and manually review upstream changes. +2. **No silent failures.** Every validation step must produce visible output. If a check can't run + (e.g., no GPU for Level 5), say so explicitly rather than skipping silently. +3. **Rollback is always an option.** If things go sideways, the user should never feel stuck. Always + have `git revert -m 1 ` ready. diff --git a/skills/te-fl-upstream-sync/phases/01-setup-and-analysis.md b/skills/te-fl-upstream-sync/phases/01-setup-and-analysis.md new file mode 100644 index 0000000000..0d4125a668 --- /dev/null +++ b/skills/te-fl-upstream-sync/phases/01-setup-and-analysis.md @@ -0,0 +1,184 @@ +### Stage 1: Repo Setup & Branch Preparation (`/stage1-setup`) + +This stage gets the repo ready: clone, add upstream remote, and create the `dev` and `base` branches +needed for the merge. You can also run sub-steps individually: `/stage1-clone`, `/stage1-create-dev`, +`/stage1-create-base`. + +#### Step 1: Clone the fork + +```bash +git clone https://github.com/flagos-ai/TransformerEngine-FL.git +cd TransformerEngine-FL +``` + +Skip if already cloned — run the Repo Detection Preamble above instead. + +#### Step 2: Add upstream remote and fetch + +```bash +git remote -v | grep upstream +# If not present: +git remote add upstream https://github.com/Nvidia/TransformerEngine.git +git fetch upstream --tags +``` + +#### Step 3: Create dev branch from upstream release (`/stage1-create-dev`) + +The `dev` branch mirrors the target upstream release exactly — no fork-specific changes. + +**Before creating the dev branch, ask the user for two parameters:** + +1. **Target upstream branch** — the upstream release to sync to (e.g. `release_v2.14`, `main`) +2. **Target upstream commit** — (optional) specific commit SHA on the target branch to checkout; + if empty, use the branch tip + +**Prompt the user:** +> Please specify: +> 1. Target upstream branch to sync to (e.g. `release_v2.14`) +> 2. Target upstream commit (leave empty for branch tip, e.g. `abc1234`) + +Store the answers: +```bash +TARGET_UPSTREAM_BRANCH="" # e.g. release_v2.14 +TARGET_UPSTREAM_COMMIT="" # e.g. abc1234 (empty = branch tip) +``` + +1. List available upstream releases for reference: + ```bash + git branch -r | grep upstream/release + ``` + +2. Create the dev branch: + ```bash + if [ -n "$TARGET_UPSTREAM_COMMIT" ]; then + git checkout -b dev ${TARGET_UPSTREAM_COMMIT} + echo "Created dev at commit ${TARGET_UPSTREAM_COMMIT} on upstream/${TARGET_UPSTREAM_BRANCH}" + else + git checkout -b dev upstream/${TARGET_UPSTREAM_BRANCH} + echo "Created dev at tip of upstream/${TARGET_UPSTREAM_BRANCH}" + fi + ``` + +3. Record the sync point — create `SYNC_POINT.md` at repo root: + ```markdown + # Upstream Sync Point + - Upstream: Nvidia/TransformerEngine + - Branch: ${TARGET_UPSTREAM_BRANCH} + - Commit SHA: + - Sync Date: + - Synced By: + ``` + +4. Verify: + ```bash + git log --oneline -5 + ``` + +#### Step 4: Create base branch from fork's original upstream (`/stage1-create-base`) + +The `base` branch represents the upstream version the fork was originally based on. This is needed +for accurate three-way merges. + +**Ask the user for two more parameters:** + +3. **Base upstream branch** — the upstream release the fork is currently based on (e.g. `release_v2.9`) +4. **Base upstream commit** — (optional) specific commit SHA on the base branch to checkout; + if empty, use the branch tip + +**Prompt the user:** +> Please specify: +> 3. Base upstream branch the fork is currently based on (e.g. `release_v2.9`) +> 4. Base upstream commit (leave empty for branch tip, e.g. `def5678`) + +Store the answers: +```bash +BASE_UPSTREAM_BRANCH="" # e.g. release_v2.9 +BASE_UPSTREAM_COMMIT="" # e.g. def5678 (empty = branch tip) +``` + +Create the base branch: +```bash +git fetch upstream ${BASE_UPSTREAM_BRANCH} + +if [ -n "$BASE_UPSTREAM_COMMIT" ]; then + git checkout -b base ${BASE_UPSTREAM_COMMIT} + echo "Created base at commit ${BASE_UPSTREAM_COMMIT} on upstream/${BASE_UPSTREAM_BRANCH}" +else + git checkout -b base upstream/${BASE_UPSTREAM_BRANCH} + echo "Created base at tip of upstream/${BASE_UPSTREAM_BRANCH}" +fi + +git log --oneline -5 +``` + +**Success criteria:** `dev` and `base` branches exist and match their respective upstream releases +commit-for-commit. + +--- + +### Stage 2: Identify Plugin Changes (`/stage2-diff-plugin-changes`) + +Before merging, you need to understand exactly what the fork added on top of upstream. This diff +between `base` (the upstream release the fork is based on) and `main` (fork) reveals all plugin-related changes — the +files you must protect during the merge. + +**Steps:** + +1. Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +2. Generate a summary of all changes the fork introduced: + ```bash + git diff base..main --stat + ``` + +3. Generate the full diff and save it for reference: + ```bash + git diff base..main > plugin_changes.diff + ``` + +4. Identify plugin-specific changes: + ```bash + # Files added or modified in plugin directory + git diff base..main --name-status -- 'transformer_engine/plugin/' + + # CUDA patches added or modified + git diff base..main --name-status -- 'transformer_engine/__init__.py' + + # Build system changes for plugin support + git diff base..main -- setup.py CMakeLists.txt pyproject.toml + + # API changes (e.g. torch.Tensor('cuda') -> torch.Tensor('TE_DEVICE_TYPE')) + # Detailed changes captured by diff base..main + git diff base..main --name-status -- 'transformer_engine/pytorch/' + ``` + +5. Record the changes — save a structured summary to `PLUGIN_CHANGES.md`: + ```markdown + # Plugin Changes (base → main) + + ## New Files (added by fork) + + + ## Modified Files (changed by fork) + + + ## Plugin Directory Contents + + + ## CUDA Patches Contents + + + ## Build System Modifications + + + ## Python Binding Modifications + + ``` + +This record is critical — during Stage 3 (Merge & Conflict Resolution), use it to verify that every +plugin change from main survives the merge. If a file listed here has a conflict, it needs +careful attention. + +**Success criteria:** `plugin_changes.diff` and `PLUGIN_CHANGES.md` generated, all fork-specific +changes catalogued. + diff --git a/skills/te-fl-upstream-sync/phases/02-merge-and-integrate.md b/skills/te-fl-upstream-sync/phases/02-merge-and-integrate.md new file mode 100644 index 0000000000..e23d0b20b2 --- /dev/null +++ b/skills/te-fl-upstream-sync/phases/02-merge-and-integrate.md @@ -0,0 +1,962 @@ +### Stage 3: Merge & Conflict Resolution (`/stage3-merge`, `/stage3-analyze-conflicts`, `/stage3-resolve-p0-conflict`, `/stage3-resolve-build-conflict`) + +This is where upstream changes meet the fork's plugin system. Conflicts are expected and normal. +This stage covers both the merge itself and resolving any conflicts that arise. + +Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +#### Step 1: Execute the merge (`/stage3-merge`) + +1. Ensure working tree is clean: + ```bash + git status + ``` + If dirty, ask the user to stash or commit first. + +2. Checkout main: + ```bash + git checkout main + ``` + +3. Create a dated merge branch (recommended for safety): + ```bash + git checkout -b merge/dev-to-main-$(date +%Y%m%d) + ``` + +4. Execute the merge with no-fast-forward: + ```bash + git merge dev --no-ff -m "merge(dev): integrate upstream release_v2.14" + ``` + +5. If the merge completes cleanly (rare but possible), skip to Stage 4. + +6. If conflicts occur, list them: + ```bash + git diff --name-only --diff-filter=U + ``` + +**Important:** Do NOT use `--strategy-option theirs` or `--strategy-option ours` globally. Each conflict +needs individual analysis based on the priority matrix. + +#### Step 2: Auto-resolve unmodified files + +Some conflicted files may not have been modified by the fork at all (i.e., `git diff base..main -- ` +is empty). For these files, it is safe to accept the upstream (dev) version directly. + +```bash +for file in $(git diff --name-only --diff-filter=U); do + if git diff base..main -- "$file" | grep -q '^'; then + echo "FORK-MODIFIED: $file" + else + echo "AUTO-RESOLVE (take upstream): $file" + git checkout --theirs "$file" + git add "$file" + fi +done +``` + +Only files printed as `FORK-MODIFIED` require priority-based resolution below. + +#### Step 3: Analyze and categorize conflicts (`/stage3-analyze-conflicts`) + +1. List all remaining conflicted files: + ```bash + git diff --name-only --diff-filter=U + ``` + +2. Categorize each file into P0/P1/P2 based on path matching: + - P0: paths containing `transformer_engine/pytorch/` + - P1: `setup.py`, `CMakeLists.txt`, `pyproject.toml`, or files in core algorithm dirs + - P2: everything else, such as `.github` cicd related + +3. Present a table to the user showing file, priority, and recommended resolution strategy. + +4. Suggest resolution order: all P0 first, then P1, then P2. + +#### Step 4: Resolve P0 conflicts (`/stage3-resolve-p0-conflict `) + +For P0 files (plugin interfaces, CUDA patches): + +1. Show the conflict diff for the specific file +2. Read the main branch version entirely. If newly add contents in main, copy to current branch. +3. Read the main branch version entirely. If `cuda` -> `te-device-type` related modification or other modification related to plugin system or patch system, apply to current branch. +3. Check security or critical bug, and fix it. + +#### Step 5: Resolve P1 build conflicts (`/stage3-resolve-build-conflict `) + +For P1 build files (setup.py, CMakeLists.txt, pyproject.toml): + +1. Show the three-way diff (base, ours, theirs) +2. Identify plugin-specific build sections in the fork version (look for comments like + `# Plugin build targets`, or targets referencing `plugin/`) +3. Identify upstream improvements (new dependencies, version bumps, new build targets) +4. Merge manually: keep all plugin build sections from main, integrate upstream changes around them +5. After resolution, verify plugin build targets still exist: + ```bash + grep -n "plugin" setup.py CMakeLists.txt + ``` + +For P2 files, use standard merge resolution — accept both sides where possible, prefer upstream +for upstream-specific content. + +#### Step 6: Finalize merge + +After all conflicts are resolved: +```bash +git add -A +pre-commit run --all-files +git add -A # re-stage any formatting fixes +git commit --no-edit +``` + +### Stage 4: Plugin API Sync (`/stage4-add-apis`, `/stage4-verify-pybind-coverage`) + +**Why this stage matters:** The main branch's plugin OP APIs were built to match the C++ bindings in +`pytorch/csrc/` as they existed in the base branch. After merging dev into main, `pytorch/csrc/` now +reflects the upstream release — which may have added new APIs, changed existing function signatures, or +removed deprecated ones. The plugin layer must be updated to stay consistent, otherwise you get +`AttributeError: module 'transformer_engine_torch' has no attribute 'xxx'` at runtime. + +**Core principle:** Plugin OP APIs must mirror `pytorch/csrc/` pybind bindings 1:1. The diff between +base and dev in `pytorch/csrc/` tells you exactly what changed; the plugin layer must reflect those +same changes. + +Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +**Key files to update:** +- `transformer_engine/plugin/core/ops.py` — base class `TEFLBackendBase` with abstract method stubs for every op +- Vendor backend implementations (one per vendor, each delegating to its own `tex` module): + - `transformer_engine/plugin/core/backends/vendor/cuda/cuda.py` — `CUDABackend`, tex = `transformer_engine_torch_nv` + - `transformer_engine/plugin/core/backends/vendor/enflame/enflame.py` — `EnflameBackend`, tex = `migration.patches.transformer_engine.v2_9_0` + - `transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py` — `IluvatarBackend`, tex = `transformer_engine_iluvatar.pytorch.ixte_torch` + - `transformer_engine/plugin/core/backends/vendor/metax/metax.py` — `MetaxBackend`, tex = `transformer_engine_torch_metax` + - `transformer_engine/plugin/core/backends/vendor/musa/musa.py` — `MUSABackend`, tex = `transformer_engine_musa_torch` + - `transformer_engine/plugin/core/backends/vendor/hygon/hygon.py` — `HygonBackend`, tex = `transformer_engine_torch_hygon` +- Vendor `register_ops.py` files (one per vendor, containing `OpImpl` registrations): + - `transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py` + - `transformer_engine/plugin/core/backends/vendor/enflame/register_ops.py` + - `transformer_engine/plugin/core/backends/vendor/iluvatar/register_ops.py` + - `transformer_engine/plugin/core/backends/vendor/metax/register_ops.py` + - `transformer_engine/plugin/core/backends/vendor/musa/register_ops.py` + - `transformer_engine/plugin/core/backends/vendor/hygon/register_ops.py` +- Scan-only (check for changed interfaces, usually no changes needed): + - `transformer_engine/plugin/core/backends/flagos/` — partial implementation, subset of ops + - `transformer_engine/plugin/core/backends/reference/` — partial implementation, subset of ops + +#### Step 1: Diff csrc between base and dev to identify all API changes + +The base branch represents the upstream version that main's plugin APIs were originally built against. +The dev branch represents the new upstream release. Diffing base..dev shows exactly what changed. + +```bash +# Full diff of C++ bindings between the old upstream (base) and new upstream (dev) +git diff base..dev -- transformer_engine/pytorch/csrc/ > /tmp/csrc_diff.diff + +echo "=== csrc diff summary ===" +diffstat /tmp/csrc_diff.diff 2>/dev/null || echo "(diffstat not available, check /tmp/csrc_diff.diff manually)" +``` + +#### Step 2: Extract and categorize API changes (ADD / MODIFY / REMOVE) + +Parse the diff to build a clear picture of what needs to change in the plugin layer. + +```bash +# --- 2a: Newly ADDED pybind APIs (lines added with .def() in the diff) --- +echo "=== ADDED APIs ===" +grep -E '^\+.*\.def\(' /tmp/csrc_diff.diff | grep -v '^\+\+\+' | sed 's/^+//' | sed 's/^[[:space:]]*//' + +# --- 2b: REMOVED pybind APIs (lines removed with .def() in the diff) --- +echo "" +echo "=== REMOVED APIs ===" +grep -E '^\-.*\.def\(' /tmp/csrc_diff.diff | grep -v '^\-\-\-' | sed 's/^-//' | sed 's/^[[:space:]]*//' + +# --- 2c: MODIFIED APIs — appear in both added and removed with the same function name --- +echo "" +echo "=== MODIFIED APIs (name appears in both added and removed lines) ===" +ADDED_NAMES=$(grep -E '^\+.*\.def\("' /tmp/csrc_diff.diff | grep -v '^\+\+\+' | \ + sed 's/.*\.def("\([^"]*\)".*/\1/' | sort -u) +REMOVED_NAMES=$(grep -E '^\-.*\.def\("' /tmp/csrc_diff.diff | grep -v '^\-\-\-' | \ + sed 's/.*\.def("\([^"]*\)".*/\1/' | sort -u) +comm -12 <(echo "$ADDED_NAMES") <(echo "$REMOVED_NAMES") + +# --- 2d: Purely new APIs (added but not in removed — truly new) --- +echo "" +echo "=== PURELY NEW APIs ===" +comm -23 <(echo "$ADDED_NAMES") <(echo "$REMOVED_NAMES") + +# --- 2e: Purely removed APIs (removed but not re-added — truly deleted) --- +echo "" +echo "=== PURELY REMOVED APIs ===" +comm -13 <(echo "$ADDED_NAMES") <(echo "$REMOVED_NAMES") + +# --- 2f: Also check Python-side references to transformer_engine_torch for new call sites --- +echo "" +echo "=== New Python-side transformer_engine_torch references ===" +git diff base..dev -- transformer_engine/pytorch/ ':(exclude)transformer_engine/pytorch/csrc/' | \ + grep -E '^\+.*transformer_engine_torch\.' | grep -v '^\+\+\+' | \ + sed 's/.*transformer_engine_torch\.//' | sed 's/[^a-zA-Z0-9_].*//' | sort -u +``` + +Review the output carefully. Save the categorized list — you will use it in the next steps. + +#### Step 2g: Detect class-object parameter changes (indirect API changes) + +Some plugin ops accept class objects (dataclasses, named tuples, or custom classes) as parameters +rather than primitive types. When the class definition changes (fields added, removed, or renamed), +the op's effective interface has changed even though its function signature is identical. The pybind +diff from Steps 2a–2f will NOT catch these — you must check explicitly. + +**Why this matters:** Consider `get_attention_backend(attention_params: AttentionParams)`. If upstream +adds new fields to `AttentionParams` (e.g., `bottom_right_diagonal`, `cuda_graph`, `num_splits`), +the function signature is unchanged, so Steps 2a–2f report no change. But vendor backends that +construct or inspect `AttentionParams` fields will break or produce wrong results if they don't +account for the new fields. This is an indirect API change that must be detected and handled. + +**How to detect:** + +1. Identify all plugin ops whose parameters include class/dataclass types. The most common cases: + - `get_attention_backend(attention_params: AttentionParams)` — `AttentionParams` is a dataclass + in `transformer_engine/pytorch/attention/dot_product_attention/utils.py` + - Any op parameter annotated with a class type defined in the TE codebase (not stdlib types) + +2. For each such class, diff its definition between base and dev: + +```bash +# --- 2g: Class-object parameter change detection --- +echo "=== Detecting class-object parameter changes ===" + +# AttentionParams — consumed by get_attention_backend +echo "" +echo "--- AttentionParams field diff (base vs dev) ---" +echo "Base fields:" +git show base:transformer_engine/pytorch/attention/dot_product_attention/utils.py 2>/dev/null | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:" | sort +echo "" +echo "Dev fields:" +git show dev:transformer_engine/pytorch/attention/dot_product_attention/utils.py 2>/dev/null | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:" | sort + +echo "" +echo "--- Fields ADDED in dev ---" +diff <(git show base:transformer_engine/pytorch/attention/dot_product_attention/utils.py 2>/dev/null | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:" | sed 's/[[:space:]]//g' | sort) \ + <(git show dev:transformer_engine/pytorch/attention/dot_product_attention/utils.py 2>/dev/null | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:" | sed 's/[[:space:]]//g' | sort) \ + | grep '^>' | sed 's/^> / NEW: /' + +echo "" +echo "--- Fields REMOVED in dev ---" +diff <(git show base:transformer_engine/pytorch/attention/dot_product_attention/utils.py 2>/dev/null | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:" | sed 's/[[:space:]]//g' | sort) \ + <(git show dev:transformer_engine/pytorch/attention/dot_product_attention/utils.py 2>/dev/null | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:" | sed 's/[[:space:]]//g' | sort) \ + | grep '^<' | sed 's/^< / REMOVED: /' + +# Repeat for any other class-object parameters found in plugin ops. +# To discover them, scan ops.py for parameter type annotations that reference TE-defined classes: +echo "" +echo "--- Other class-typed parameters in TEFLBackendBase ops ---" +grep -E "def \w+\(self.*:\s*[A-Z]\w+Params|def \w+\(self.*:\s*Optional\[[A-Z]\w+Params\]" \ + transformer_engine/plugin/core/ops.py | grep -v "FlashAttentionBase" +``` + +3. If new fields are found, treat the consuming op as MODIFIED — even though its signature is + unchanged. Add it to the MODIFIED list from Step 2c. Then in Steps 4–5, update the plugin's + handling of that op: + - If the vendor backend passes the class object through transparently (e.g., CUDA/MetaX/MUSA + just forward `attention_params` to the upstream implementation), no code change is needed — + the new fields flow through automatically. Document this. + - If the vendor backend constructs the class object, inspects its fields, or has custom logic + that ignores the parameter (e.g., Hygon/Iluvatar/Reference use env-var logic and ignore + `attention_params`), verify that the new fields don't break the custom logic. If the vendor + ignores the parameter entirely, no change is needed — but document the finding. + - If the class field change affects the op's return value or behavior in a way that downstream + code depends on, the vendor implementation may need updating. + +#### Step 3: Cross-reference with current plugin ops + +Before making changes, check what already exists in the plugin layer to understand the gap. + +```bash +# List all ops currently defined in the base class +echo "=== Current plugin op definitions (ops.py base class) ===" +grep -n "def " transformer_engine/plugin/core/ops.py | head -80 + +# List all backend implementations in the CUDA reference backend +echo "" +echo "=== Current CUDA backend implementations ===" +grep -n "def " transformer_engine/plugin/core/backends/vendor/cuda/cuda.py | head -80 + +# List all registered ops in the CUDA register_ops.py +echo "" +echo "=== Current CUDA registered ops ===" +grep "op_name" transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py | head -80 + +# Cross-reference: find APIs from Step 2 that are NOT yet in the plugin +echo "" +echo "=== Missing from plugin (need to ADD) ===" +for api_name in $ADDED_NAMES; do + if ! grep -q "$api_name" transformer_engine/plugin/core/ops.py 2>/dev/null; then + echo " MISSING: $api_name" + fi +done + +echo "" +echo "=== In plugin but signature may need UPDATE ===" +for api_name in $(comm -12 <(echo "$ADDED_NAMES") <(echo "$REMOVED_NAMES")); do + if grep -q "$api_name" transformer_engine/plugin/core/ops.py 2>/dev/null; then + echo " CHECK: $api_name (exists in plugin, signature may have changed)" + fi +done +``` + +This gives you a clear action list: which APIs to add, which to update, and which to remove. + +Use the CUDA backend as the reference implementation — it is always the most complete and should be +updated first. All other vendor backends mirror CUDA's method signatures. + +#### Step 4: Update base class op definitions (ops.py) + +For each API change identified in Steps 2-3, update `transformer_engine/plugin/core/ops.py` (the +`TEFLBackendBase` class): + +- **For ADDED APIs:** Add a new abstract method stub mirroring the C++ function's Python-visible + signature. Read the full `.def(...)` block in the new csrc code to get the exact parameter names + and types. Follow the pattern of existing methods in the class — each method is a thin stub that + raises `NotImplementedError` or returns a default value. + +- **For MODIFIED APIs:** Update the existing method signature to match the new parameters. Compare + the old and new `.def(...)` blocks side by side. Common changes: added/removed parameters, changed + default values, renamed arguments. + +- **For REMOVED APIs:** If an API was removed from csrc, decide whether to: + - Remove it from the plugin (if nothing in the fork depends on it) + - Keep it with a deprecation warning (if fork-specific code still uses it) + - Mark it clearly with a comment for manual review + +- **CommOverlapP2P special case:** If new APIs relate to `CommOverlapP2P`, check whether they should + be class methods on the `CommOverlapP2P` wrapper in ops.py (which routes to backend + `create_comm_overlap_p2p`) or standalone methods on `TEFLBackendBase`. + +#### Step 5: Update CUDA reference backend first (cuda.py + register_ops.py) + +The CUDA backend is the reference implementation. Update it first, then replicate to other vendors. + +**5a: Update `backends/vendor/cuda/cuda.py`** + +Mirror every change from Step 4 in the `CUDABackend` class: + +- **For ADDED ops:** Add a method that delegates to `tex.xxx(...)` where `tex` is the vendor's + C extension module (`transformer_engine_torch_nv` for CUDA). Match the parameter list exactly + with the ops.py base class definition. Key patterns to follow: + - DType conversion: use `tex.DType(int(dtype))` when passing TE dtype enums to the C extension + - Tensor parameters: pass through directly + - Return values: return whatever `tex.xxx()` returns + +- **For MODIFIED ops:** Update the method to pass the new parameters correctly. If parameters were + added/removed, the `tex.xxx()` call must reflect that. + +- **For REMOVED ops:** Remove or deprecate the corresponding method. + +**5b: Update `backends/vendor/cuda/register_ops.py`** + +For each ADDED op, add an `OpImpl` registration entry. Follow the existing pattern: + +```python +OpImpl( + op_name="new_api_name", + impl_id="vendor.cuda", + vendor="NVIDIA", + priority=100, + impl=CUDABackend.new_api_name, + is_available=_bind_is_available(CUDABackend.new_api_name), +), +``` + +The `_bind_is_available` helper wraps the backend method for `OpImpl.is_available()` checks. +Place new entries in the same logical grouping as related existing ops. + +#### Step 5c: Replicate to all other vendor backends + +After CUDA is complete and verified, replicate the same changes to all other vendor backends. +Each vendor follows the identical pattern — only the class name, `tex` module, `impl_id`, and +`vendor` string differ: + +| Vendor | Class | tex module | impl_id | vendor string | +|----------|------------------|------------------------------------------------|--------------------|---------------| +| CUDA | `CUDABackend` | `transformer_engine_torch_nv` | `vendor.cuda` | `NVIDIA` | +| Enflame | `EnflameBackend` | `migration.patches.transformer_engine.v2_9_0` | `vendor.enflame` | `ENFLAME` | +| Iluvatar | `IluvatarBackend`| `transformer_engine_iluvatar.pytorch.ixte_torch`| `vendor.iluvatar` | `Iluvatar` | +| MetaX | `MetaxBackend` | `transformer_engine_torch_metax` | `vendor.metax` | `METAX` | +| MUSA | `MUSABackend` | `transformer_engine_musa_torch` | `vendor.musa` | `MUSA` | +| Hygon | `HygonBackend` | `transformer_engine_torch_hygon` | `vendor.hygon` | `HYGON` | + +For each vendor, update both files: +1. `{vendor}/{vendor}.py` — add/modify/remove the same methods as CUDA, delegating to that + vendor's `tex` module +2. `{vendor}/register_ops.py` — add/modify/remove the same `OpImpl` entries with the vendor's + `impl_id` and `vendor` string + +**Efficiency tip:** Since all vendors have identical method bodies (only `tex` module differs), +you can write a batch script to apply the same changes across all vendors simultaneously rather +than editing each file manually. Diff the CUDA register_ops op_names against each vendor's to +identify exactly which ops are missing. + +**Note on Enflame:** Enflame's `tex` module is `migration.patches.transformer_engine.v2_9_0` +(loaded lazily via `_get_tex()`). It also has a custom `flash_attention.py` and +`get_attention_backend` implementation that routes through its own migration layer. When adding +new ops, follow the same delegation pattern as other vendors. + +**Note on Hygon:** Hygon may have a pre-existing gap in OpImpl count compared to other vendors +(e.g., 8 fewer ops). This is expected — only add the new ops from this sync, don't try to +backfill the pre-existing gap. + +**Critical: Explicit parameter signatures required.** All vendor backend methods MUST use +explicit parameter lists matching the CUDA backend exactly. Do NOT use `*args, **kwargs` as a +shortcut — this hides interface mismatches and causes silent failures when upstream adds new +parameters. The only exceptions are methods where CUDA itself uses `*args, **kwargs` (e.g., +`te_general_grouped_gemm_for_*`, `nvfp4_compute_per_block_scale`, `nvfp4_expand_scale_to_fp8`, +`nvfp4_fused_scale`, `nvfp4_multi_tensor_2d_partial_cast`). + +Example — correct: +```python +def mxfp8_scaling_compute_partial_amax( + self, + tensor: torch.Tensor, + amax: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, +) -> None: + tex = self._get_tex() + return tex.mxfp8_scaling_compute_partial_amax(tensor, amax, h, w, start_offset, block_len) +``` + +Example — wrong (do not do this): +```python +def mxfp8_scaling_compute_partial_amax(self, *args, **kwargs): + tex = self._get_tex() + return tex.mxfp8_scaling_compute_partial_amax(*args, **kwargs) +``` + +After adding/modifying methods in all vendors, verify consistency: +```bash +# Verify no new *args/**kwargs were introduced (excluding known exceptions) +for vendor in enflame hygon iluvatar metax musa; do + f="transformer_engine/plugin/core/backends/vendor/$vendor/$vendor.py" + echo "=== $vendor ===" + grep -n "def.*\*args.*\*\*kwargs" "$f" | \ + grep -v "te_general_grouped_gemm_for_\|nvfp4_compute_per_block_scale\|nvfp4_expand_scale_to_fp8\|nvfp4_fused_scale\|nvfp4_multi_tensor_2d_partial_cast" +done +# If any output appears, those methods need explicit parameter lists from CUDA. +``` + +#### Step 5d: Scan flagos and reference backends for changed interfaces + +The `flagos/` and `reference/` backends are partial implementations that only cover a subset of +ops. They do NOT need new method stubs for newly added APIs. However, if any MODIFIED APIs +(signature changes) affect methods that these backends implement, those signatures must be updated. + +```bash +# Check if any of the modified API names exist in flagos/reference backends +echo "=== Checking flagos backend for modified APIs ===" +for api_name in ; do + grep -rn "def $api_name" transformer_engine/plugin/core/backends/flagos/ 2>/dev/null +done + +echo "" +echo "=== Checking reference backend for modified APIs ===" +for api_name in ; do + grep -rn "def $api_name" transformer_engine/plugin/core/backends/reference/ 2>/dev/null +done +``` + +If any matches are found, update those method signatures to match the new parameters from Step 4. +If no matches are found (the common case), no changes are needed — document this in the commit log. + +#### Step 5e: Add plugin-to-tex type conversions for new methods + +The plugin layer (ops.py) defines its own Python enum types (`DType`, `NVTE_QKV_Layout`, +`NVTE_Bias_Type`, `NVTE_Mask_Type`, `NVTE_Softmax_Type`, `NVTE_QKV_Format`, `CommOverlapType`, +etc.) that mirror the C++ enums but are distinct Python objects. Each vendor's `tex` module has +its own versions of these enums (e.g., `tex.DType`, `tex.NVTE_QKV_Layout`). Since they share +the same integer values but are different types, conversion is required at the boundary when +calling `tex.*()`. + +After adding all new methods to vendor backends, audit each one for parameters that carry +plugin-layer enum types or objects containing such types. Two conversion patterns apply: + +**Pattern 1 — Standalone enum parameters:** +When a method receives a bare enum parameter (e.g., `dtype: DType`, `otype: DType`, +`qkv_layout`, `bias_type`, `attn_mask_type`), convert it before passing to `tex`: + +```python +dtype = tex.DType(int(dtype)) if dtype is not None else None +qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None +bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None +comm_type = tex.CommOverlapType(int(comm_type)) if comm_type is not None else None +``` + +The general form is `tex.(int(param))` with a `None` guard when the parameter is +optional. Look at the pybind `.def()` signature to determine which parameters are enums — they +will have C++ types like `DType`, `NVTE_QKV_Layout`, etc. + +**Pattern 2 — Object attribute conversion (quantizer pattern):** +When a method receives a complex object (like a `quantizer`) that internally holds enum-typed +attributes, those attributes must be normalized before the object is passed to `tex`. The +quantizer's `.dtype` attribute is the most common case: + +```python +# Normalize quantizer.dtype to this backend's `tex.DType`. +try: + if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"): + qdtype = quantizer.dtype + if qdtype is not None: + quantizer.dtype = tex.DType(int(qdtype)) +except Exception: + pass +``` + +This pattern is defensive (`try/except`, `hasattr` checks) because the quantizer object comes +from the Python layer and its structure may vary. Currently applies to `quantize` and +`bgrad_quantize`, but any new API that accepts a quantizer or similar wrapper object needs the +same treatment. + +**How to identify which new methods need conversion:** +1. Read each new method's pybind `.def()` signature in csrc +2. Any parameter typed as a TE enum (`DType`, `NVTE_*`, `CommOverlapType`, etc.) → Pattern 1 +3. Any parameter that is a complex object containing TE enums (e.g., quantizer) → Pattern 2 +4. Plain tensors, ints, floats, bools, strings → no conversion needed + +This step applies to all vendor backends equally — the conversion logic is identical, only the +`tex` module differs. + +#### Step 6: Verify API consistency across all backends + +Before committing, verify that the plugin layer is fully consistent with csrc across all vendors. + +```bash +echo "=== Verification: Plugin API Consistency Check ===" + +# 6a: Extract all pybind API names from the current (merged) csrc +ALL_CSRC_APIS=$(grep -rh '\.def("' transformer_engine/pytorch/csrc/ 2>/dev/null | \ + sed 's/.*\.def("\([^"]*\)".*/\1/' | sort -u) + +# 6b: Extract all method names from the base class (ops.py) +ALL_OPS_METHODS=$(grep -E "^\s+def " transformer_engine/plugin/core/ops.py 2>/dev/null | \ + sed 's/.*def \([a-zA-Z_][a-zA-Z0-9_]*\).*/\1/' | grep -v "^__" | sort -u) + +# 6c: Find csrc APIs missing from ops.py base class +echo "--- csrc APIs missing from ops.py (should be empty) ---" +comm -23 <(echo "$ALL_CSRC_APIS") <(echo "$ALL_OPS_METHODS") + +# 6d: Verify all vendor backends have the same method count as CUDA reference +echo "" +echo "=== Vendor backend method counts (should match CUDA) ===" +for vendor in cuda enflame iluvatar metax musa hygon; do + VENDOR_FILE="transformer_engine/plugin/core/backends/vendor/$vendor/$vendor.py" + if [ -f "$VENDOR_FILE" ]; then + COUNT=$(grep -c "^\s*def " "$VENDOR_FILE" 2>/dev/null) + echo " $vendor: $COUNT methods" + else + echo " $vendor: FILE NOT FOUND" + fi +done + +# 6e: Verify all vendor register_ops have consistent OpImpl counts +echo "" +echo "=== Vendor register_ops OpImpl counts ===" +for vendor in cuda enflame iluvatar metax musa hygon; do + REG_FILE="transformer_engine/plugin/core/backends/vendor/$vendor/register_ops.py" + if [ -f "$REG_FILE" ]; then + COUNT=$(grep -c "op_name" "$REG_FILE" 2>/dev/null) + echo " $vendor: $COUNT op_name entries" + else + echo " $vendor: FILE NOT FOUND" + fi +done + +# 6f: Quick sanity check +CSRC_COUNT=$(echo "$ALL_CSRC_APIS" | wc -l | tr -d ' ') +OPS_COUNT=$(echo "$ALL_OPS_METHODS" | wc -l | tr -d ' ') +echo "" +echo "csrc API count: $CSRC_COUNT" +echo "ops.py method count: $OPS_COUNT" + +if [ "$CSRC_COUNT" -le "$OPS_COUNT" ]; then + echo "✅ Plugin covers all csrc APIs" +else + echo "⚠️ Plugin may be missing $(($CSRC_COUNT - $OPS_COUNT)) API(s) — review the list above" +fi + +# 6g: Syntax-check all modified Python files +echo "" +echo "=== Syntax validation ===" +for f in transformer_engine/plugin/core/ops.py \ + transformer_engine/plugin/core/backends/vendor/cuda/cuda.py \ + transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py \ + transformer_engine/plugin/core/backends/vendor/enflame/enflame.py \ + transformer_engine/plugin/core/backends/vendor/enflame/register_ops.py \ + transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py \ + transformer_engine/plugin/core/backends/vendor/iluvatar/register_ops.py \ + transformer_engine/plugin/core/backends/vendor/metax/metax.py \ + transformer_engine/plugin/core/backends/vendor/metax/register_ops.py \ + transformer_engine/plugin/core/backends/vendor/musa/musa.py \ + transformer_engine/plugin/core/backends/vendor/musa/register_ops.py \ + transformer_engine/plugin/core/backends/vendor/hygon/hygon.py \ + transformer_engine/plugin/core/backends/vendor/hygon/register_ops.py; do + python3 -c "import ast; ast.parse(open('$f').read())" 2>&1 && echo " ✅ $f" || echo " ❌ $f" +done +``` + +If the verification shows missing APIs, go back to Steps 4-5 and fix them before proceeding. + +#### Step 7: Sync FlashAttention Class Hierarchy with Upstream + +Steps 1–6 above handle `TEFLBackendBase` ops (including ordinary methods like `get_attention_backend` +that accept class-object parameters — see Step 2g for how those are detected). This step handles a +separate concern: the **FlashAttention class hierarchy**, which is independent of `TEFLBackendBase` +and uses its own inheritance-based plugin pattern. + +**Architecture:** The plugin defines `FlashAttentionBase` in `ops.py` (inherits `torch.nn.Module` + +`ABC`). Each vendor provides a subclass that either delegates to a vendor-specific flash attention +library or implements its own attention logic: + +- **Delegation pattern** (CUDA, MetaX, MUSA, Hygon, Iluvatar): The vendor subclass's `_forward_impl` + has the same signature as `FlashAttentionBase._forward_impl` and delegates the actual computation + to a vendor library. When upstream adds a new parameter, these subclasses need the parameter added + to their signature and passed through to the delegation call. +- **Custom implementation pattern** (KunlunXin, Reference, FlagOS): The vendor subclass implements + its own attention logic (e.g., using `torch.nn.functional.scaled_dot_product_attention` or + `flag_gems`). When upstream adds a new parameter, these subclasses need the parameter in their + signature but may choose to ignore it or implement support for it depending on their backend's + capabilities. + +**Note on `get_attention_backend` and `FusedAttention`:** `get_attention_backend` is an ordinary +method inside `TEFLBackendBase` — it is handled by Steps 1–6 (with class-object parameter changes +like `AttentionParams` detected by Step 2g). `FusedAttention` is an upstream class that the plugin +does NOT wrap (no `FusedAttentionBase` exists) — upstream's class is used directly, so signature +changes to `FusedAttention.forward()` do not require plugin changes. + +**Files to update:** +- `transformer_engine/plugin/core/ops.py` — `FlashAttentionBase._forward_impl` + `forward` signatures +- `transformer_engine/plugin/core/backends/vendor/cuda/flash_attention.py` +- `transformer_engine/plugin/core/backends/vendor/enflame/flash_attention.py` +- `transformer_engine/plugin/core/backends/vendor/hygon/flash_attention.py` +- `transformer_engine/plugin/core/backends/vendor/metax/flash_attention.py` +- `transformer_engine/plugin/core/backends/vendor/musa/flash_attention.py` +- `transformer_engine/plugin/core/backends/vendor/iluvatar/flash_attention.py` (if exists) +- `transformer_engine/plugin/core/backends/vendor/kunlunxin/flash_attention.py` (if exists) +- `transformer_engine/plugin/core/backends/reference/flash_attention.py` +- `transformer_engine/plugin/core/backends/flagos/attention/dot_product_attention/backends.py` + +**7a: Diff upstream FlashAttention.forward() signature changes** + +```bash +# --- FlashAttention.forward() --- +echo "=== Base branch FlashAttention.forward signature ===" +git show base:transformer_engine/pytorch/attention/dot_product_attention/backends.py 2>/dev/null | \ + sed -n '/class FlashAttention/,/"""flash-attn fprop"""/p' | grep -E "^\s+\w+.*[:=]" + +echo "" +echo "=== Dev branch (current) FlashAttention.forward signature ===" +sed -n '/class FlashAttention(torch.nn.Module)/,/"""flash-attn fprop"""/p' \ + transformer_engine/pytorch/attention/dot_product_attention/backends.py | grep -E "^\s+\w+.*[:=]" +``` + +For each interface, identify parameters that were added, removed, or had their defaults changed. + +For `get_attention_backend`: the function signature may be unchanged while `AttentionParams` gains +new fields — check both. Also check whether the function body reads new fields from +`AttentionParams` (compare the field extraction block at the top of the function). + +For `FusedAttention`: if the plugin does not wrap it (no `FusedAttentionBase` in `ops.py`), then +signature changes do not require plugin updates — document this finding and move on. + +Report the diff results. If no parameters were added/removed/changed, FlashAttention needs no plugin +updates — document this and skip to Step 8. If changes exist, proceed with 7b–7e below. + +**7b: Compare with plugin FlashAttentionBase** + +```bash +echo "=== Plugin FlashAttentionBase._forward_impl signature ===" +sed -n '/class FlashAttentionBase/,/raise NotImplementedError/p' \ + transformer_engine/plugin/core/ops.py | grep -E "^\s+\w+.*[:=]" + +echo "" +echo "=== Plugin FlashAttentionBase.forward signature ===" +sed -n '/class FlashAttentionBase/,/def backend_name/p' \ + transformer_engine/plugin/core/ops.py | grep -E "^\s+\w+.*[:=]" +``` + +Cross-reference each parameter with the upstream `FlashAttention.forward()` signature from 7a. +Any parameter present in upstream but missing from the plugin is a gap that must be filled. + +**7c: Update FlashAttentionBase in ops.py** + +For each new/changed parameter identified in 7a–7b: + +1. Add the parameter to `_forward_impl()` abstract method signature (with the same default value as upstream) +2. Add the parameter to `forward()` method signature (same default) +3. Add the parameter to the `_forward_impl()` call inside `forward()` (the direct call path) +4. Add the parameter to the `call_impl_fn` lambda/closure inside `forward()` (the fallback dispatch path) + +Preserve the parameter ordering from upstream. New parameters typically go at the end, before +`**kwargs` if present. + +**7d: Update all vendor FlashAttention subclasses** + +For each vendor's `flash_attention.py`: + +1. Add the new parameter(s) to `_forward_impl()` with the same default value +2. For **delegation-pattern** vendors (CUDA, MetaX, MUSA, Hygon, Iluvatar): pass the new parameter + through to the delegation call (e.g., the `tex.fused_attn_*` call or the upstream function call) +3. For **custom-implementation** vendors (KunlunXin, Reference, FlagOS): add the parameter to the + signature. Whether to implement support depends on the backend's capabilities: + - If the parameter controls an optimization the backend doesn't support (e.g., `num_splits` for + a torch SDPA backend), accept the parameter but don't use it — the default value should be + safe to ignore + - If the parameter changes semantics (e.g., a new attention mask type), the implementation may + need updating + +```bash +# Check all vendor flash_attention files for current _forward_impl signatures +echo "=== Vendor FlashAttention _forward_impl signatures ===" +for vendor in cuda enflame hygon metax musa iluvatar kunlunxin; do + FILE="transformer_engine/plugin/core/backends/vendor/$vendor/flash_attention.py" + if [ -f "$FILE" ]; then + echo "--- $vendor ---" + grep -A 30 "def _forward_impl" "$FILE" | head -35 | grep -E "^\s+\w+.*[:=]" + echo "" + fi +done + +echo "--- reference ---" +grep -A 30 "def _forward_impl" \ + transformer_engine/plugin/core/backends/reference/flash_attention.py | head -35 | grep -E "^\s+\w+.*[:=]" + +echo "" +echo "--- flagos ---" +grep -A 30 "def _forward_impl" \ + transformer_engine/plugin/core/backends/flagos/attention/dot_product_attention/backends.py | \ + head -35 | grep -E "^\s+\w+.*[:=]" +``` + +Update each file to include the new parameter(s). For delegation-pattern vendors, also update the +delegation call to pass the new parameter through. + +**7e: Update flagos and reference backends (custom implementations)** + +The flagos (`FlashAttentionFL`) and reference (`FlashAttentionTorch`) backends have their own +attention implementations rather than delegating to upstream. When new parameters are added: + +1. Add the parameter to `_forward_impl()` signature (same default as upstream) +2. Decide whether to implement support: + - Parameters that affect the core attention computation (e.g., new mask types, new dropout + modes) should be implemented if the backend supports them, or raise `NotImplementedError` + with a clear message if not + - Parameters that are optimization hints (e.g., `num_splits` for controlling parallelism) + can safely be accepted and ignored — the default value produces correct results +3. For KunlunXin's `FlashAttentionTorch` (uses `torch.nn.functional.scaled_dot_product_attention`): + check if PyTorch's SDPA supports the new parameter natively + +```bash +# Verify upstream call sites pass the new parameter(s) +echo "=== Upstream flash_attention call sites ===" +grep -n "self\.flash_attention(" transformer_engine/pytorch/ -r --include="*.py" -A 30 | \ + grep -E "self\.flash_attention\(|^\s+\w+\s*=" | head -10 + +echo "" +echo "=== Plugin FlashAttentionBase params (after update) ===" +grep -A 30 "def _forward_impl" transformer_engine/plugin/core/ops.py | \ + head -35 | grep -E "^\s+\w+.*[:=]" +``` + +Confirm every argument passed at each upstream call site is accepted by the plugin's +`FlashAttentionBase._forward_impl` and all vendor subclass implementations. + +#### Step 8: Commit and log all changes + +```bash +# Build a detailed change log +LOG_FILE="/tmp/plugin_api_changes.log" +echo "=== Plugin API changes for upstream sync ===" > "$LOG_FILE" +echo "Date: $(date -Iseconds)" >> "$LOG_FILE" +echo "Diff base: base..dev" >> "$LOG_FILE" +echo "" >> "$LOG_FILE" + +echo "--- Files changed ---" >> "$LOG_FILE" +git diff --name-only HEAD >> "$LOG_FILE" +echo "" >> "$LOG_FILE" + +echo "--- Op definition changes (ops.py base class) ---" >> "$LOG_FILE" +git diff HEAD -- transformer_engine/plugin/core/ops.py >> "$LOG_FILE" +echo "" >> "$LOG_FILE" + +echo "--- Vendor backend changes ---" >> "$LOG_FILE" +for vendor in cuda enflame iluvatar metax musa hygon; do + echo "=== $vendor ===" >> "$LOG_FILE" + git diff HEAD -- "transformer_engine/plugin/core/backends/vendor/$vendor/" >> "$LOG_FILE" + echo "" >> "$LOG_FILE" +done + +cat "$LOG_FILE" + +git add -A +pre-commit run --all-files +git add -A # re-stage any formatting fixes +git commit -m "plugin: sync plugin APIs with upstream csrc changes + +Updated plugin OP API layer to match pytorch/csrc/ pybind changes +between base and dev branches. Changes applied to: +- ops.py base class (TEFLBackendBase) +- ops.py FlashAttentionBase (synced forward/\_forward\_impl signatures with upstream FlashAttention) +- All vendor FlashAttention subclasses (cuda, enflame, hygon, metax, musa, iluvatar, kunlunxin) +- All 6 vendor backends (cuda, enflame, iluvatar, metax, musa, hygon) +- All 6 vendor register_ops.py files +- Scanned flagos/reference backends for changed interfaces +See /tmp/plugin_api_changes.log for details." +``` + + +#### Step 9: Full-Surface Pybind Coverage Audit (`/stage4-verify-pybind-coverage`) + +Steps 1–8 above handle new/modified/removed APIs found by diffing `pytorch/csrc/` between base and dev. +But omissions can also come from: +- Pybind exports in `common/util/pybind_helper.h` (enums, utility functions via `NVTE_DECLARE_COMMON_PYBIND11_HANDLES`) +- Functions added in earlier upstream versions that were never wrapped in the plugin layer +- Exports that the diff-based steps missed (e.g., unchanged but unwrapped) + +This step does a full-surface audit: extract every pybind-exported symbol, compare against the plugin +layer, and flag any gaps. + +**Key source files:** +- `transformer_engine/pytorch/csrc/extensions/pybind.cpp` — main `PYBIND11_MODULE` with `m.def()` calls +- `transformer_engine/common/util/pybind_helper.h` — `NVTE_DECLARE_COMMON_PYBIND11_HANDLES` macro (enums + utility functions) +- `transformer_engine/plugin/core/ops.py` — `TEFLBackendBase` class (abstract method stubs for every op) +- `transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py` — `OpImpl` registrations + +#### Step 1: Extract all pybind-exported function names + +```bash +# 1a: Extract m.def() function names from pybind.cpp +grep -oP 'm\.def\("\K[^"]+' transformer_engine/pytorch/csrc/extensions/pybind.cpp | sort -u > /tmp/pybind_exports.txt + +# 1b: Extract utility function names from pybind_helper.h +# These are inside NVTE_DECLARE_COMMON_PYBIND11_HANDLES macro +grep -oP 'm\.def\("\K[^"]+' transformer_engine/common/util/pybind_helper.h >> /tmp/pybind_exports.txt + +# 1c: Check for any other pybind modules in common/ or pytorch/csrc/ +grep -rl 'PYBIND11_MODULE\|m\.def(' transformer_engine/common/ transformer_engine/pytorch/csrc/ 2>/dev/null | \ + grep -v pybind.cpp | grep -v pybind_helper.h + +sort -u -o /tmp/pybind_exports.txt /tmp/pybind_exports.txt +echo "Total pybind exports: $(wc -l < /tmp/pybind_exports.txt)" +``` + +#### Step 2: Extract all plugin-layer op names + +```bash +# 2a: Method names from TEFLBackendBase in ops.py +grep -oP 'def \K\w+' transformer_engine/plugin/core/ops.py | grep -v '^_' | sort -u > /tmp/plugin_ops.txt + +# 2b: Registered op_names from CUDA register_ops.py (most complete vendor) +grep -oP 'op_name="\K[^"]+' transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py | sort -u > /tmp/registered_ops.txt + +echo "TEFLBackendBase methods: $(wc -l < /tmp/plugin_ops.txt)" +echo "CUDA registered ops: $(wc -l < /tmp/registered_ops.txt)" +``` + +#### Step 3: Find omissions + +```bash +# 3a: Pybind exports missing from TEFLBackendBase +echo "=== Pybind exports NOT in TEFLBackendBase ===" +comm -23 /tmp/pybind_exports.txt /tmp/plugin_ops.txt + +# 3b: Pybind exports missing from CUDA register_ops.py +echo "" +echo "=== Pybind exports NOT in CUDA register_ops.py ===" +comm -23 /tmp/pybind_exports.txt /tmp/registered_ops.txt + +# 3c: In TEFLBackendBase but not registered (abstract method with no implementation) +echo "" +echo "=== In TEFLBackendBase but NOT registered in CUDA ===" +comm -23 /tmp/plugin_ops.txt /tmp/registered_ops.txt +``` + +#### Step 4: Categorize and triage omissions + +For each omission found in Step 3, categorize it: + +| Category | Action | +|----------|--------| +| Compute ops (kernels, gemm, attention, norm, activation) | Must add: ops.py stub + vendor impl + register_ops.py | +| Utility/query functions (device_supports_X, get_version, etc.) | May skip if only used internally by the C++ layer, or add if Python code calls `tex.xxx()` | +| Enum/class exports (DType, FP8TensorMeta, etc.) | Handled separately by TEFLModule constructor — verify they're in the enum/class setup | + +For each compute op omission, check if Python code actually calls it: +```bash +# For each missing op, check if it's called via tex.xxx +for op in $(comm -23 /tmp/pybind_exports.txt /tmp/registered_ops.txt); do + count=$(grep -r "tex\.$op" transformer_engine/pytorch/ --include="*.py" 2>/dev/null | wc -l) + if [ "$count" -gt 0 ]; then + echo "CRITICAL: tex.$op called $count times but not in plugin layer" + else + echo "LOW: $op not called via tex in Python code" + fi +done +``` + +#### Step 5: Fix omissions (same pattern as Steps 4–5 above) + +For each CRITICAL omission (pybind export that Python code calls via `tex.xxx` but plugin doesn't wrap): + +1. Add abstract method stub to `TEFLBackendBase` in `ops.py`, matching the pybind signature +2. Add implementation to `CUDABackend` in `cuda.py` that delegates to `self._get_tex().xxx(...)` +3. Add `OpImpl` registration in `cuda/register_ops.py` +4. Repeat for ALL other vendor backends (iluvatar, metax, musa, hygon) — each vendor must get: + - A method in its backend class (e.g., `iluvatar.py`, `metax.py`, `musa.py`, `hygon.py`) + that delegates to its own native tex module via `self._get_tex().xxx(...)` + - An `OpImpl` registration in its `register_ops.py` + - If the vendor's native module doesn't support the function, implement a safe fallback + (e.g., return `False` for query functions, raise `NotImplementedError` for compute ops) +5. Check if flagos/reference backends need an implementation + +For LOW-priority omissions (not called via `tex` in Python), document them but don't add unless needed. + +#### Step 6: Verify completeness + +```bash +# Re-run the comparison to confirm zero critical omissions remain +echo "=== Remaining omissions ===" +comm -23 /tmp/pybind_exports.txt /tmp/registered_ops.txt | while read op; do + count=$(grep -r "tex\.$op" transformer_engine/pytorch/ --include="*.py" 2>/dev/null | wc -l) + if [ "$count" -gt 0 ]; then + echo "STILL MISSING: tex.$op ($count call sites)" + fi +done +echo "If no output above, all critical omissions are resolved." +``` + +#### Pre-commit and commit + +Run pre-commit hooks before committing to ensure code formatting is correct: + +```bash +pre-commit run --files $(git diff --name-only --cached) || true +# If pre-commit modified files, stage them again +git add -u +``` + +``` +git commit -m "fix: add missing pybind exports to plugin layer + +Audited all pybind-exported symbols against plugin ops.py / register_ops.py. +Found N omissions, M critical (called via tex.xxx in Python code). +Added ops.py stubs, vendor implementations, and registrations for: +- +Skipped N low-priority utility functions not called from Python." +``` + diff --git a/skills/te-fl-upstream-sync/phases/03-patch-and-verify.md b/skills/te-fl-upstream-sync/phases/03-patch-and-verify.md new file mode 100644 index 0000000000..8c3ea2e5c2 --- /dev/null +++ b/skills/te-fl-upstream-sync/phases/03-patch-and-verify.md @@ -0,0 +1,432 @@ +### Stage 5: Patch CUDA Hardcoding in Upstream Python Changes (`/stage5-patch-cuda-hardcoding`) + +Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +The main branch carries a "patch" feature that abstracts CUDA-specific string references in +TransformerEngine's Python layer. The key mechanism: + +- `TE_DEVICE_TYPE` (string, default `"cuda"`) — defined in `transformer_engine/__init__.py`, + overridden by vendor `patches.py` at import time (e.g., MUSA sets it to `"musa"`) + +Only `TE_DEVICE_TYPE` is needed in this stage. The `torch.cuda.*` API calls (streams, events, +synchronize, device queries, etc.) are handled separately by vendor `patches.py` files via +runtime monkey-patching (e.g., `torch.cuda.synchronize → torch.musa.synchronize`). Those +calls stay as `torch.cuda.*` in source code — no replacement needed. + +When upstream merges introduce new or modified Python files, they may bring fresh `"cuda"` +string hardcoding. This stage detects and patches those instances. + +#### What to patch vs what to leave alone + +**Patch** — hardcoded `"cuda"` strings used as device identifiers: + +| Pattern | Replacement | +|---------|-------------| +| `device="cuda"` | `device=TE_DEVICE_TYPE` | +| `torch.device("cuda")` | `torch.device(TE_DEVICE_TYPE)` | +| `torch.get_autocast_dtype("cuda")` | `torch.get_autocast_dtype(TE_DEVICE_TYPE)` | +| `.device.type == "cuda"` | `.device.type == TE_DEVICE_TYPE` | +| `"cuda"` in other device-selection contexts | `TE_DEVICE_TYPE` | + +**Leave as-is** — everything else: + +- All `torch.cuda.*` API calls — handled by vendor patches.py at runtime +- `torch.cuda.CUDAGraph` — CUDA-specific, no vendor equivalent +- `torch.cuda.nvtx.*` — profiling, handled by patches.py or no-op +- `torch.version.cuda` — build-time version query +- `.cuda_stream` — low-level C pointer +- `"cuda"` in comments, docstrings, and string messages +- `"cuda"` in device-type checks that guard CUDA-specific code blocks (these are intentional + gates, not device selection) + +#### Step 1: Scan the upstream Python diff for new `"cuda"` string hardcoding + +```bash +git diff base..dev -- transformer_engine/pytorch/ ':(exclude)transformer_engine/pytorch/csrc/' \ + > /tmp/python_layer_diff.diff +``` + +Extract newly added lines containing `"cuda"` string patterns: + +```bash +grep '^+' /tmp/python_layer_diff.diff | grep -v '^+++' \ + | grep -E 'device.*"cuda"|torch\.device\("cuda"\)|get_autocast_dtype.*"cuda"|\.device\.type.*==.*"cuda"' \ + > /tmp/cuda_string_candidates.txt +``` + +This is the candidate list. Each line needs manual triage in Step 2. + +#### Step 2: Triage candidates + +For each candidate line, decide patch or skip: + +1. **Patch** if the `"cuda"` string is used for device selection in general-purpose code + (modules, ops, distributed, quantization, attention, etc.) +2. **Skip** if the `"cuda"` string is inside a CUDA-specific guard (e.g., + `if device.type == "cuda": `) — these are intentional gates +3. **Skip** if it's in a docstring, comment, or log message +4. **Skip** if the file is inherently CUDA-only (e.g., `cuda_graphs.py`) + +The key distinction: device *selection* (`device="cuda"`) should use `TE_DEVICE_TYPE` so +non-CUDA vendors get their device. Device *detection* (`if x == "cuda"`) that gates +CUDA-specific behavior should stay as `"cuda"`. + +#### Step 3: Apply patches + +For each file with lines to patch: + +1. Read the current file content (post-merge, on the working branch) +2. Replace `"cuda"` string patterns per the table above +3. Add the import if not already present: + ```python + from transformer_engine import TE_DEVICE_TYPE + ``` +4. Syntax check: `python3 -c "import ast; ast.parse(open('').read())"` + +**Important**: Only patch lines introduced or modified by the upstream merge (the `^+` lines +from the diff). Do not retroactively patch pre-existing `"cuda"` references that the main +branch has already chosen to leave as-is. + +**Do not modify** any files under `plugin/core/backends/vendor/` — those are vendor-customized. + +#### Step 4: Verify + +1. Syntax check all modified files +2. Grep to confirm no un-patched `"cuda"` device strings remain in newly added lines: + ```bash + git diff base..HEAD -- transformer_engine/pytorch/ ':(exclude)transformer_engine/pytorch/csrc/' \ + | grep '^+' | grep -v '^+++' \ + | grep -E 'device.*=.*"cuda"|torch\.device\("cuda"\)|get_autocast_dtype\("cuda"\)' \ + > /tmp/remaining_cuda_strings.txt + ``` + Review — each remaining line should be a deliberate skip (guard, docstring, CUDA-only file). +3. Count patched vs skipped for the commit log. + +#### Step 5: Commit + +```bash +git add -A +pre-commit run --all-files +git add -A # re-stage any formatting fixes +git commit -m "patch: normalize new upstream 'cuda' string hardcoding to TE_DEVICE_TYPE + +Scanned Python-layer diff (base..dev, excluding csrc) for newly introduced +hardcoded 'cuda' device strings. Replaced instances across files: +- device='cuda' → device=TE_DEVICE_TYPE: +- torch.device('cuda') → torch.device(TE_DEVICE_TYPE): +- get_autocast_dtype('cuda') → get_autocast_dtype(TE_DEVICE_TYPE): +- .device.type == 'cuda' → .device.type == TE_DEVICE_TYPE: +Skipped intentional guards and CUDA-specific blocks. +torch.cuda.* API calls left as-is (handled by vendor patches.py at runtime)." +``` + + +### Stage 6: Detect & Fix Stale References in Fork-Specific Code (`/stage6-detect-stale-refs`) + +Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +After Stages 7 and 8, the merge branch contains both upstream updates (from dev) and fork-specific +additions (from main). Because main was originally based on the base branch, fork-specific code may +reference functions, classes, or file paths that upstream has since renamed or relocated between base +and dev. These stale references won't cause merge conflicts (the fork code is new relative to dev), +but they will cause runtime errors — calling a function that no longer exists or importing from a +module that has moved. + +**The core problem:** The merge branch = main + dev. Code that is new in the merge branch compared +to dev (i.e., fork-specific code) was written against the base branch's API surface. If dev renamed +`_load_cudnn` to `_load_cudnn_v2`, or moved `quantized_tensor.py` to `float8/quantized_tensor.py`, +the fork code still references the old names. + +#### Step 1: Identify what upstream renamed or moved between base and dev + +Build two inventories of upstream changes: + +**1a. Renamed/removed Python symbols (functions, classes, constants):** + +```bash +# Get all Python files that changed between base and dev (upstream evolution) +git diff --name-only base..dev -- '*.py' > /tmp/upstream_changed_py_files.txt + +# Extract removed/renamed function and class definitions +# Lines starting with '-' that define functions or classes (but not '---' diff headers) +git diff base..dev -- '*.py' | \ + grep -E '^\-[^-]' | \ + grep -E '^\-(def |class | def )' | \ + sed 's/^-//' | \ + sed 's/(.*//' | \ + sed 's/def //' | sed 's/class //' | \ + sed 's/://' | \ + tr -d ' ' | \ + sort -u > /tmp/upstream_removed_symbols.txt + +# Extract added/renamed function and class definitions +git diff base..dev -- '*.py' | \ + grep -E '^\+[^+]' | \ + grep -E '^\+(def |class | def )' | \ + sed 's/^+//' | \ + sed 's/(.*//' | \ + sed 's/def //' | sed 's/class //' | \ + sed 's/://' | \ + tr -d ' ' | \ + sort -u > /tmp/upstream_added_symbols.txt + +# Symbols that were removed but NOT re-added are truly gone +# Symbols removed AND re-added with a different name are renames +comm -23 /tmp/upstream_removed_symbols.txt /tmp/upstream_added_symbols.txt \ + > /tmp/upstream_gone_symbols.txt + +echo "=== Symbols removed/renamed in upstream (base→dev) ===" +cat /tmp/upstream_gone_symbols.txt +echo "Count: $(wc -l < /tmp/upstream_gone_symbols.txt)" +``` + +**1b. Relocated files (moved or renamed):** + +```bash +# Detect file renames/moves between base and dev +git diff --diff-filter=R --name-status -M base..dev > /tmp/upstream_renamed_files.txt + +# Also detect deleted files (might have been moved without git detecting the rename) +git diff --diff-filter=D --name-only base..dev > /tmp/upstream_deleted_files.txt + +echo "=== Files renamed/moved in upstream ===" +cat /tmp/upstream_renamed_files.txt +echo "" +echo "=== Files deleted in upstream ===" +cat /tmp/upstream_deleted_files.txt +``` + +If both lists are empty, this stage is a no-op — skip to the commit step. + +#### Step 2: Identify fork-specific code (new in merge branch vs dev) + +This is the code at risk — it was written against the base branch API and has never been reconciled +with upstream's renames. + +```bash +# Fork-specific content = lines present in current branch (merge result) but not in dev +# Focus on Python files in plugin/ and any fork-specific directories +git diff dev..HEAD -- '*.py' | \ + grep -E '^\+[^+]' | \ + grep -v '^+++' \ + > /tmp/fork_new_lines.txt + +# Also get the list of fork-specific files (files that exist in HEAD but not in dev) +git diff --name-only dev..HEAD -- '*.py' > /tmp/fork_changed_files.txt + +echo "=== Fork-specific changed files ===" +cat /tmp/fork_changed_files.txt +echo "Count: $(wc -l < /tmp/fork_changed_files.txt)" +``` + +#### Step 3: Cross-reference — find stale references + +For each gone/renamed symbol from Step 1, check if fork-specific code references it: + +```bash +echo "=== Scanning fork-specific code for stale references ===" +STALE_FOUND=0 + +while IFS= read -r symbol; do + [ -z "$symbol" ] && continue + # Search fork-changed files for references to this symbol + MATCHES=$(grep -rn "$symbol" --include='*.py' \ + $(cat /tmp/fork_changed_files.txt) 2>/dev/null | \ + grep -v "^Binary" || true) + if [ -n "$MATCHES" ]; then + echo "" + echo "⚠️ STALE REFERENCE: '$symbol' (removed/renamed in upstream)" + echo "$MATCHES" + STALE_FOUND=$((STALE_FOUND + 1)) + fi +done < /tmp/upstream_gone_symbols.txt + +# Check for imports from relocated/deleted files +while IFS= read -r old_file; do + [ -z "$old_file" ] && continue + # Convert file path to module path for import matching + MODULE=$(echo "$old_file" | sed 's/\.py$//' | sed 's/\//./g') + BASENAME=$(basename "$old_file" .py) + MATCHES=$(grep -rn "import.*$BASENAME\|from.*$MODULE" --include='*.py' \ + $(cat /tmp/fork_changed_files.txt) 2>/dev/null | \ + grep -v "^Binary" || true) + if [ -n "$MATCHES" ]; then + echo "" + echo "⚠️ STALE IMPORT: references deleted/moved file '$old_file'" + echo "$MATCHES" + STALE_FOUND=$((STALE_FOUND + 1)) + fi +done < /tmp/upstream_deleted_files.txt + +echo "" +echo "=== Summary: $STALE_FOUND stale reference(s) found ===" +``` + +#### Step 4: Resolve stale references + +For each stale reference found, determine the correct replacement: + +1. **Renamed symbol:** Find the new name in dev by searching for the function's context: + ```bash + # Example: if _load_cudnn was renamed, find what replaced it + git log --all --oneline --diff-filter=M base..dev -- + git diff base..dev -- | grep -A5 -B5 "old_symbol_name" + ``` + The diff context usually shows the old name removed and the new name added nearby. + +2. **Relocated file:** Use the rename detection from Step 1b, or search dev for the file: + ```bash + # Find where the file moved to + git ls-tree -r --name-only dev | grep "" + ``` + +3. **Truly removed (no replacement):** The upstream removed the functionality entirely. The fork + code needs to be refactored to use the replacement API, or the old implementation needs to be + kept as a fork-specific utility. Flag these for manual review. + +Apply each fix, keeping fork-specific logic intact while updating references to match the current +upstream API surface. + +#### Step 5: Verify fixes + +```bash +# Re-run the stale reference scan — should find 0 issues +# (repeat Step 3 commands) + +# Syntax-check all modified files +for f in $(git diff --name-only HEAD -- '*.py'); do + python3 -c "import ast; ast.parse(open('$f').read())" 2>&1 && \ + echo " ✅ $f" || echo " ❌ $f" +done +``` + +#### Step 6: Commit + +```bash +git add -A +pre-commit run --all-files +git add -A # re-stage any formatting fixes +git commit -m "fix: update stale references in fork code to match upstream renames + +Scanned fork-specific code (new in merge vs dev) for references to +functions, classes, and file paths that upstream renamed or relocated +between base and dev. Fixed stale reference(s): +- +- " +``` + + +### Stage 7: Build & Import Verification (`/stage7-basic-test`) + +This stage validates that the merged code compiles and the core import chain works. It is the +first gate after all code changes (Stages 3–6). If this fails, nothing else matters — fix it +before moving on. + +Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +#### Step 1: Environment Setup & Build + +```bash +# Ensure we are in the TransformerEngine-FL directory (Repo Detection Preamble) +if [ -d "TransformerEngine-FL" ]; then + cd TransformerEngine-FL +elif [ "$(basename $(pwd))" = "TransformerEngine-FL" ]; then + echo "Already in TransformerEngine-FL" +elif [ -d "../TransformerEngine-FL" ]; then + cd ../TransformerEngine-FL +else + echo "ERROR: TransformerEngine-FL directory not found. Run /stage1-setup first." + exit 1 +fi +echo "Working directory: $(pwd)" + +# Check and display current conda environment +conda info --envs +echo "Active conda env: $CONDA_DEFAULT_ENV" +conda info +python --version +which python +which pip + +# Update third-party dependencies (upstream may bump submodule versions) +git submodule update --init --recursive + +# NOTE: --no-build-isolation is REQUIRED. Without it, pip creates an isolated venv +# that won't have access to the current conda env's PyTorch/CUDA dependencies. +# NOTE: Do NOT remove build/ directories or .so files before installing. +# The editable install handles incremental builds correctly. +pip install --no-build-isolation -e . 2>&1 | tee build.log + +# Verify shared objects were generated +find . -name "*.so" -newer build.log | head -10 +``` + +#### Step 2: Import Verification + +```bash +python -c "from transformer_engine import pytorch" +# Expected output: +# [CUDA] Successfully loaded CUDA libs +# [TE-FL manager.py INFO] OpManager initialized: 110 ops with 179 implementations +# [TE-FL manager.py INFO] Registered impl_ids: ['default.flagos', 'reference.torch', 'vendor.cuda'] +``` + +If the import succeeds, the basic test passes — report success and proceed to Stage 8. + +#### Step 3: Failure Diagnosis & Iterative Fix + +If the import test fails (`from transformer_engine import pytorch`), do NOT proceed. Instead, +diagnose which of the previous stages (3/4/5/6) caused the issue and fix it. + +```bash +# 3a: Trace the failing import chain +python -c " +import traceback +try: + from transformer_engine import pytorch +except Exception: + traceback.print_exc() +" + +# 3b: Common failure patterns — analyze the traceback to identify the root cause: +# +# - "transformer_engine_torch.xxx" AttributeError +# → Upstream added new C++ APIs that the fork's plugin backend doesn't expose yet. +# → Fix: Re-run Stage 4 (Plugin API Sync) to add the missing APIs. +# +# - ImportError / ModuleNotFoundError for a moved/renamed module +# → Upstream renamed or moved files between releases. +# → Detect: git diff --name-status base..dev | grep '^R' (renamed) +# git diff --name-status base..dev | grep '^D' (deleted/moved) +# → Fix: Update fork imports to the new upstream path. This is typically +# a Stage 3 (conflict resolution) or Stage 5 (CUDA patching) issue. +# +# - SyntaxError or merge conflict markers in source +# → A Stage 3 conflict resolution left bad content. +# → Fix: Go back to Stage 3 and re-resolve the affected file. + +# 3c: Cross-branch investigation to pinpoint the source +# git show dev: > /tmp/file_dev.py +# git show main: > /tmp/file_main.py +# git show base: > /tmp/file_base.py +# diff /tmp/file_base.py /tmp/file_dev.py # what upstream changed +# diff /tmp/file_base.py /tmp/file_main.py # what fork changed +# diff /tmp/file_dev.py /tmp/file_main.py # divergence between fork and upstream +# This reveals whether the bug is from a bad merge resolution, an upstream +# rename the fork still references, or a fork addition conflicting with upstream. + +# 3d: Decision — fix directly or rollback +# If the problem can be directly fixed (e.g., missing import, typo, small patch): +# → Fix it, commit with descriptive message, re-run Step 2. +# If the problem is systemic (e.g., entire stage needs re-execution): +# → Rollback to the commit before that stage, re-execute the stage with +# updated steps, then re-run from Step 1. + +# 3e: After fixing, re-verify +python -c "from transformer_engine import pytorch; print('OK')" +python -c "from transformer_engine import te_device_type; print(te_device_type())" +``` + +Iterate until the import succeeds. Each fix should be committed separately with a descriptive message +(e.g., `fix: resolve import error for X after upstream merge`). + diff --git a/skills/te-fl-upstream-sync/phases/04-test-and-finalize.md b/skills/te-fl-upstream-sync/phases/04-test-and-finalize.md new file mode 100644 index 0000000000..f661a2267c --- /dev/null +++ b/skills/te-fl-upstream-sync/phases/04-test-and-finalize.md @@ -0,0 +1,443 @@ +### Stage 8: Unit & Integration Tests (`/stage8-run-unit-tests`, `/stage8-run-full-cicd`) + +After Stage 7 confirms the build and import work, run the full test suite. Three levels of testing, +each building on the previous. If any level fails, stop and diagnose before proceeding. + +Run the Repo Detection Preamble to ensure you are in the TransformerEngine-FL directory. + +#### Pre-check: CI Script Validation + +Before running any tests, verify that the CI scripts in `qa/` reference test files that actually +exist. Upstream renames test files between releases (e.g. `test_float8tensor.py` → +`test_quantized_tensor.py`), and the TE-FL CI scripts may not have been updated to match. + +Run this check: + +```bash +# Extract all .py test file references from qa/L0_pytorch_unittest/test.sh and verify they exist +grep -oP '(?<=\$TE_PATH/)tests/pytorch/[^\s"]+\.py' qa/L0_pytorch_unittest/test.sh | sort -u | while read f; do + if [ ! -f "$f" ]; then + echo "MISSING: $f" + fi +done +# Also check directory references (e.g. nvfp4/) +grep -oP '(?<=\$TE_PATH/)tests/pytorch/[^\s"]+(? | head -5 + ``` +2. Update the CI script to use the new filename. +3. Also check whether new test files exist in `tests/pytorch/` that are not yet referenced in the + CI script — compare against the upstream version of `qa/L0_pytorch_unittest/test.sh`: + ```bash + git show upstream/main:qa/L0_pytorch_unittest/test.sh | grep -oP 'tests/pytorch/[^\s"]+\.py' | sort > /tmp/upstream_tests.txt + grep -oP '(?<=\$TE_PATH/)tests/pytorch/[^\s"]+\.py' qa/L0_pytorch_unittest/test.sh | sort > /tmp/local_tests.txt + diff /tmp/upstream_tests.txt /tmp/local_tests.txt + ``` + Lines prefixed with `<` are in upstream but missing locally — add them if the files exist. + +4. Fix any issues, run `pre-commit run --files qa/L0_pytorch_unittest/test.sh`, and commit before + proceeding to the test levels below. + +**Also check for tests that are in the skip list but never actually run** — a test appearing only +in the MetaX/CUDA skip block but not in any `run_test_step` call is a sign it was added to the +skip list when the test was introduced upstream, but the `run_test_step` call was never added. + +**Known patterns to watch for after each upstream sync:** +- Test file renames (check `git log --diff-filter=R` on the upstream merge commit) +- New test files added under `tests/pytorch/` or `tests/pytorch/attention/` +- Tests moved into subdirectories (e.g. `test_attention.py` → `attention/test_attention.py`) + +#### Per-Level Results Table + +After each level finishes, parse the test output and display a results table to the console. +This gives immediate visibility into what passed/failed before moving to the next level. + +The per-level table format (one row per test function): + +| # | Test Name | Result | Duration | Details | +|---|-----------|--------|----------|---------| +| 1 | test_foo | ✅ PASSED | 0.3s | | +| 2 | test_bar | ❌ FAILED | 1.2s | AssertionError: expected X | +| 3 | test_baz | ⚠️ SKIPPED | — | reason: no GPU | +| | **Total** | **2/3 passed** | **1.5s** | **1 failed, 1 skipped** | + +How to compile the table: +- Parse the pytest `-v` output: each line like `test_file.py::test_name PASSED/FAILED/SKIPPED` is a row +- Extract duration from the pytest summary line (e.g., `=== 42 passed in 5.23s ===`) +- For FAILED tests, include the first line of the failure reason from the `--tb=short` traceback +- For SKIPPED tests, include the skip reason if available +- Add a summary row at the bottom with totals + +#### Level 1 — Plugin-Specific Tests +```bash +pytest transformer_engine/plugin/tests/ -k "plugin" -v --tb=short 2>&1 | tee plugin-test.log +``` + +**→ Display Level 1 results table, then continue.** + +#### Level 2 — Integration Tests +```bash +pytest tests/pytorch/ -v --tb=short 2>&1 | tee integration-test.log +``` + +**→ Display Level 2 results table, then continue.** + +#### Level 2.5 — CI Test Suites (OP API Validation) + +This level runs three CI test suites in sequence: debug → unit → distributed. These validate the +plugin OP API interfaces across all vendor backends. This is critical after upstream merges because +upstream may change function signatures in ways that break the plugin dispatch layer. + +```bash +cd "$TE_FL_DIR" +# L0 Debug tests +TE_PATH=$(pwd) bash qa/L0_pytorch_debug_unittest/test.sh 2>&1 | tee l0-debug.log +# L0 Unit tests +TE_PATH=$(pwd) bash qa/L0_pytorch_unittest/test.sh 2>&1 | tee l0-unittest.log +# L1 Distributed tests +TE_PATH=$(pwd) bash qa/L1_pytorch_distributed_unittest/test.sh 2>&1 | tee l1-distributed.log +``` + +**Common failure pattern: OP API signature mismatch** + +If tests fail with errors like: +- `CUDABackend.fused_topk_with_score_function_bwd() takes 10 positional arguments but 11 were given` +- `CUDABackend.fused_score_for_moe_aux_loss_bwd() got an unexpected keyword argument 'grad_logits'` + +This indicates that the upstream pytorch layer changed the call signature, but the plugin OP API +definition and backend implementations were not updated to match. + +**Fix procedure:** + +1. **Identify the failing OP** from the error message (e.g., `fused_topk_with_score_function_bwd`) + +2. **Check the upstream caller** to see the expected signature: + ```bash + grep -n "tex.fused_topk_with_score_function_bwd" transformer_engine/pytorch/router.py + ``` + Note all parameters being passed, including output tensors like `grad_logits`. + +3. **Check the C++ extension signature** to confirm the expected interface: + ```bash + grep -A 10 "void fused_topk_with_score_function_bwd" transformer_engine/pytorch/csrc/extensions/router.cpp + ``` + +4. **Update the OP API abstract method** in `transformer_engine/plugin/core/ops.py`: + - Add any missing parameters (e.g., `grad_logits: torch.Tensor`) + - Ensure parameter order matches the upstream caller + +5. **Update ALL backend implementations** (not just CUDA): + ```bash + # Find all backends that implement this OP + find transformer_engine/plugin/core/backends -name "*.py" | xargs grep -l "def fused_topk_with_score_function_bwd" + ``` + + For each backend file found: + - Add the missing parameter to the function signature + - Pass it through to the underlying `tex.*` call + + **Backends to check:** + - `vendor/cuda/cuda.py` + - `vendor/hygon/hygon.py` + - `vendor/metax/metax.py` + - `vendor/enflame/enflame.py` + - `vendor/iluvatar/iluvatar.py` + - `vendor/musa/musa.py` + - `flagos/flagos.py` (if the OP is implemented) + - `reference/reference.py` (if the OP is implemented) + +6. **Run pre-commit** on all modified files: + ```bash + pre-commit run --files transformer_engine/plugin/core/ops.py \ + transformer_engine/plugin/core/backends/vendor/cuda/cuda.py \ + transformer_engine/plugin/core/backends/vendor/hygon/hygon.py \ + transformer_engine/plugin/core/backends/vendor/metax/metax.py \ + transformer_engine/plugin/core/backends/vendor/enflame/enflame.py \ + transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py \ + transformer_engine/plugin/core/backends/vendor/musa/musa.py + ``` + +7. **Commit the fix**: + ```bash + git add transformer_engine/plugin/core/ops.py transformer_engine/plugin/core/backends/vendor/*/ + git commit -m "fix(plugin): update OP API signatures for + + Upstream changed the call signature to include . + Updated abstract method in ops.py and all vendor backend implementations." + ``` + +8. **Re-run all three CI test suites** to verify the fix: + ```bash + TE_PATH=$(pwd) bash qa/L0_pytorch_debug_unittest/test.sh 2>&1 | tee l0-debug-rerun.log + TE_PATH=$(pwd) bash qa/L0_pytorch_unittest/test.sh 2>&1 | tee l0-unittest-rerun.log + TE_PATH=$(pwd) bash qa/L1_pytorch_distributed_unittest/test.sh 2>&1 | tee l1-distributed-rerun.log + ``` + +**→ Display Level 2.5 results table, then continue.** + +#### Level 2.6 — Plugin Unit Tests (`run_all_tests.py`) + +Run the plugin's own test suite, which validates backend registration, op dispatch, and +plugin-layer correctness independently of the upstream pytorch tests. + +```bash +cd "$TE_FL_DIR" +python transformer_engine/plugin/tests/run_all_tests.py 2>&1 | tee plugin-run-all.log +``` + +If tests fail, check for: +- Missing op registrations in `transformer_engine/plugin/core/register_ops.py` +- Backend method not implemented (raises `NotImplementedError`) +- Import errors caused by renamed or removed upstream symbols + +**→ Display Level 2.6 results table, then continue.** + +#### Level 3 — End-to-End Tests +```bash +python tests/pytorch/test_sanity.py 2>&1 | tee e2e-test.log +``` +For non-pytest scripts, parse stdout for pass/fail indicators and display a simplified table. + +**→ Display Level 3 results table.** + +#### Final Summary + +After all levels complete (or on first failure), display the cumulative summary table: + +| Level | Test Suite | Status | Passed | Failed | Skipped | Duration | +|-------|-----------|--------|--------|--------|---------|----------| +| L1 | Plugin Tests | ✅/❌ | N | N | N | Xs | +| L2 | Integration | ✅/❌ | N | N | N | Xs | +| L2.5a | L0 Debug Unit Tests | ✅/❌ | N | N | N | Xs | +| L2.5b | L0 PyTorch Unit Tests | ✅/❌ | N | N | N | Xs | +| L2.5c | L1 Distributed Tests | ✅/❌ | N | N | N | Xs | +| L2.6 | Plugin Unit Tests | ✅/❌ | N | N | N | Xs | +| L3 | End-to-End | ✅/❌ | N | N | N | Xs | +| | **Total** | | **N** | **N** | **N** | **Xs** | + +If any level failed, list the specific failing tests below the summary table for quick reference. + +Document the failure in the sync report so the next attempt can address it. + +#### Bug Fix Commit + +If any test failures are caused by plugin-layer issues (e.g., missing parameters, signature mismatches +between upstream callers and plugin wrappers), fix them immediately. After fixing: + +1. **Run pre-commit** on all changed files: + ```bash + pre-commit run --files ... + ``` + If pre-commit modifies files (e.g., black reformatting), re-run to confirm clean. + +2. **Commit the fix** with a descriptive message: + ```bash + git add + git commit -m "fix: + +
" + ``` + +3. **Re-run the failing test** to verify the fix before proceeding to the next level. + +Remember: fixes should cover ALL vendor backends (cuda, enflame, hygon, iluvatar, metax, musa), not just the +one being tested. Check ops.py (abstract method) + all 6 vendor backend files. + +--- + +### Stage 9: Merge to main (Tree Replacement Strategy) (`/stage9-merge-to-main`) + +When the dev/merge branch (e.g. `merge/dev-to-main-20260410`) is a **superset** of main — meaning +all main's features have been incorporated during the preceding stages — use the tree replacement +strategy to create a clean merge commit suitable for a PR to main. + +**Prerequisite:** The merge branch must be complete and verified (Stage 1–8 passed). If the merge +branch is missing features that exist on main, go back and fix it first — do NOT patch during +this stage. + +#### Why tree replacement instead of `-X theirs`? + +`-X theirs` resolves **conflicts** by taking theirs, but **non-conflicting changes from both +sides are still merged**. When both branches independently added the same patch (e.g. CUDA +patches, plugin additions), git sees them as non-conflicting additions and keeps both copies — +resulting in duplicate imports, duplicate code blocks, and broken code. Tree replacement avoids +this entirely. + +#### Step 1: Identify the source branch + +Ask the user to confirm the merge branch name. Default is `merge/dev-to-main-20260410`: + +``` +Which branch should be merged into main? +Default: merge/dev-to-main-20260410 +``` + +Store the branch name: +```bash +MERGE_BRANCH="merge/dev-to-main-20260410" # or user-provided value +``` + +#### Step 2: Create merge branch with tree replacement + +```bash +# Ensure we're in the TransformerEngine-FL directory +cd "$TE_FL_DIR" + +# Start from main +git checkout main +git pull origin main + +# Create a new branch for the PR +git checkout -b merge-to-main-$(date +%Y%m%d) + +# Tree replacement merge: +# "merge -s ours" records both parents but keeps main's tree, +# then "read-tree" replaces the tree with the merge branch's content. +git merge -s ours ${MERGE_BRANCH} --no-edit +git read-tree -m -u ${MERGE_BRANCH} + +# Run pre-commit before finalizing +pre-commit run --all-files +# If pre-commit modified files, stage them +git add -A +git commit --amend --no-edit +``` + +After this, the working tree is **identical** to `${MERGE_BRANCH}`, but the commit has both +`main` and `${MERGE_BRANCH}` as parents (preserving full history). + +#### Step 3: Remove intermediate sync records + +The merge branch may contain files created during the sync process that should not land on main +(e.g., `SYNC_POINT.md`). Remove them: + +```bash +# Remove intermediate sync/version record files +for f in SYNC_POINT.md MERGE_RECORD.md UPSTREAM_SYNC.md; do + if [ -f "$f" ]; then + git rm "$f" + fi +done + +# Commit if anything was removed +if ! git diff --cached --quiet; then + git commit -m "chore: remove intermediate sync record files (not needed on main)" +fi +``` + +#### Step 4: Verify tree equality + +```bash +# Should produce no output (or only the removed sync files) +git diff ${MERGE_BRANCH} HEAD --stat +``` + +If there is unexpected diff output beyond the removed sync files, something went wrong. +Investigate before proceeding. + +#### Step 5: Incorporate new main commits (if any) + +If `origin/main` received new commits after the merge branch was created: + +```bash +git checkout main +git pull origin main +git checkout merge-to-main-$(date +%Y%m%d) +git merge main --no-edit +# Resolve any conflicts (typically few, since merge branch is a superset) +pre-commit run --all-files +git add -A +git commit # if conflicts were resolved or pre-commit modified files +``` + +#### Step 6: Final verification + +```bash +# No conflict markers +grep -rn "<<<<<<" transformer_engine/ tests/ .github/ 2>/dev/null | head + +# History is correct — both parents present +git log --oneline --graph -5 + +# Build still works +pip install -e . 2>&1 | tail -5 +python -c "import transformer_engine; print('OK')" +``` + +**Commit discipline:** Run `pre-commit run --all-files` before every commit in this stage. +If pre-commit modifies files, re-stage and re-run until clean. + +**Checkpoint:** +1. Verify `git diff ${MERGE_BRANCH} HEAD` produces no output +2. Verify `git log --oneline --graph -3` shows both parents +3. Verify `pip install -e .` and import succeed +4. Verify `pre-commit run --all-files` passes clean + +**Success criteria:** Merge branch tree equals `${MERGE_BRANCH}`, no duplicate code blocks, both +parents in history, pip-installable, pre-commit clean. PR is submitted manually by the user. + +--- + +### Stage 10: FlagScale End-to-End Training Validation (`/stage10-flagscale-training`) + +After the merge to main is prepared (Stage 9), validate the merged TransformerEngine-FL in a real +training scenario using FlagScale. This catches runtime integration issues that unit tests miss — +wrong tensor shapes, device mismatches, plugin dispatch failures under real workloads, etc. + +**This stage delegates to the `e2e-stage-manager` skill.** Use it to: + +1. Create a new unified stage config (e.g., stage11) with the correct model × implementation × backend matrix +2. Run batch training tests across all combinations +3. Compare results with previous stages + +Refer to [`e2e-stage-manager/SKILL.md`](../e2e-stage-manager/SKILL.md) for full instructions on +config generation, batch execution, and result comparison. + +#### Sync-specific prerequisites before invoking e2e-stage-manager + +1. **Install the latest TE-FL** into the conda environment from the merge branch: + ```bash + cd /path/to/TransformerEngine-FL + pip install -e . --no-build-isolation + python -c "import transformer_engine; print(transformer_engine.__version__)" + ``` + +2. **Run pylint** to catch lint errors before training (CI runs pylint with `set -e`): + ```bash + python3 -m pylint --recursive=y transformer_engine/common transformer_engine/pytorch transformer_engine/debug + ``` + Fix any issues, run `pre-commit run --files `, and commit before proceeding. + +3. **Use a separate FlagScale installation** — do NOT use the FlagScale in the same workspace as + TransformerEngine-FL. The user must provide an absolute path to a compatible FlagScale repo. + +4. **Run on the `merge-to-main-YYYYMMDD` branch** produced by Stage 9, not the intermediate merge branch. + +#### Success criteria + +At least one parameter combination completes 20 training steps without errors, and loss values are +decreasing. The batch comparison table (generated by e2e-stage-manager) makes it easy to spot which +combinations work and which need further investigation. + +#### Fix-and-retry during training + +If a combination fails, diagnose using these common patterns: + +| Error Type | Root Cause | Fix | +|-----------|-----------|-----| +| TypeError / missing positional argument | Plugin API signature mismatch | Fix in ops.py + all vendor backends | +| AttributeError | Stale reference to renamed symbol | Find new name via `git diff base..dev` | +| RuntimeError: device mismatch | Un-patched `"cuda"` hardcoding | Replace with `TE_DEVICE_TYPE` | +| Plugin dispatch error | Missing op registration | Add to register_ops.py for all vendors | + +After each fix: pre-commit, commit, rebuild (`pip install -e . --no-build-isolation`), then rerun +the same combination to verify. Apply fixes to ALL vendor backends (cuda, enflame, musa, iluvatar, hygon, metax). diff --git a/skills/te-fl-upstream-sync/references/cicd-pipeline.md b/skills/te-fl-upstream-sync/references/cicd-pipeline.md new file mode 100644 index 0000000000..3145e5eb5a --- /dev/null +++ b/skills/te-fl-upstream-sync/references/cicd-pipeline.md @@ -0,0 +1,258 @@ +# CI/CD Pipeline Reference + +Detailed instructions for each verification level after an upstream sync merge. + +## Level 1 — Compile Test (`/run-level1-compile`) + +The most basic check: does the code compile with plugin support? + +```bash +# Clean any previous build artifacts +rm -rf build/ dist/ *.egg-info + +# Install in editable mode with verbose output to catch warnings +pip install -e . -v 2>&1 | tee build.log + +# Verify shared libraries were generated +echo "=== Checking .so files ===" +find . -name "*.so" -newer build.log | head -20 + +# Specifically check for plugin-related .so +echo "=== Plugin .so files ===" +find . -name "*plugin*" -name "*.so" + +# Check for compilation warnings related to plugin +echo "=== Plugin-related warnings ===" +grep -i "plugin\|cuda_patch" build.log | grep -i "warn\|error" || echo "No plugin warnings found" +``` + +**Pass criteria:** +- Exit code 0 from pip install +- At least one `.so` file generated +- No errors mentioning plugin or cuda_patch in build log + +**Common failures after upstream sync:** +- Missing include paths for plugin headers → check CMakeLists.txt +- Undefined symbols → check if upstream renamed/removed functions the plugin depends on +- CUDA version mismatch → check CUDA toolkit version vs upstream requirements + +--- + +## Level 2 — Unit Tests + +```bash +# Run upstream unit tests +pytest tests/unit/ -v --tb=short 2>&1 | tee unit-test.log + +# Summary +echo "=== Unit Test Summary ===" +tail -5 unit-test.log +``` + +**Pass criteria:** +- All pre-existing tests pass +- No new test failures compared to upstream's test results for the same release + +**If tests fail:** +- Check if failures are in plugin-related tests → likely a merge issue +- Check if failures are in upstream tests → might be environment issue, compare with upstream CI + +--- + +## Level 3 — Integration Tests (PyTorch) + +```bash +# Run PyTorch integration tests +pytest tests/pytorch/ -v --tb=short -x 2>&1 | tee integration-test.log + +# Summary +echo "=== Integration Test Summary ===" +tail -10 integration-test.log +``` + +**Pass criteria:** +- All PyTorch integration tests pass +- No regressions from previous sync + +**Note:** These tests may require a GPU. If running in a CPU-only environment, skip and note +in the sync report that Level 3 was not verified. + +--- + +## Level 4 — Plugin-Specific Tests (`/run-level4-plugin`) + +These are the fork-specific validations. They verify that the plugin system survived the merge intact. + +### 4.1 Plugin Directory Integrity + +```bash +echo "=== Plugin Directory Check ===" + +# Check directory exists +if [ -d "transformer_engine/common/plugin" ]; then + echo "✅ Plugin directory exists" + ls -la transformer_engine/common/plugin/ +else + echo "❌ Plugin directory MISSING" + exit 1 +fi + +# Check CUDA patches directory +if [ -d "transformer_engine/common/cuda_patches" ]; then + echo "✅ CUDA patches directory exists" + ls -la transformer_engine/common/cuda_patches/ +else + echo "❌ CUDA patches directory MISSING" + exit 1 +fi +``` + +### 4.2 OP API Interface Verification + +```bash +echo "=== OP API Interface Check ===" + +# Check register_plugin signature +grep -n "register_plugin" transformer_engine/common/plugin/*.h +if [ $? -eq 0 ]; then + echo "✅ register_plugin() found" +else + echo "❌ register_plugin() MISSING" + exit 1 +fi + +# Check unregister_plugin signature +grep -n "unregister_plugin" transformer_engine/common/plugin/*.h +if [ $? -eq 0 ]; then + echo "✅ unregister_plugin() found" +else + echo "❌ unregister_plugin() MISSING" + exit 1 +fi + +# Check for ABI-breaking changes (compare signatures with known-good) +echo "=== Signature Diff ===" +grep -A2 "register_plugin\|unregister_plugin" transformer_engine/common/plugin/*.h +``` + +### 4.3 CUDA Patches Applicability + +```bash +echo "=== CUDA Patches Check ===" + +PATCH_DIR="transformer_engine/common/cuda_patches" +FAIL=0 + +for patch in "$PATCH_DIR"/*.patch; do + [ -f "$patch" ] || continue + echo "Testing: $patch" + if git apply --check "$patch" 2>/dev/null; then + echo " ✅ Applies cleanly" + else + echo " ❌ FAILS to apply" + FAIL=1 + fi +done + +if [ $FAIL -eq 1 ]; then + echo "⚠️ Some patches need rebasing" +else + echo "✅ All patches apply cleanly" +fi +``` + +### 4.4 Build Configuration Check + +```bash +echo "=== Build Config Check ===" + +# Check setup.py for plugin references +if grep -q "plugin" setup.py 2>/dev/null; then + echo "✅ setup.py contains plugin build targets" +else + echo "⚠️ setup.py may be missing plugin targets" +fi + +# Check CMakeLists.txt for plugin targets +if grep -q "plugin" CMakeLists.txt 2>/dev/null; then + echo "✅ CMakeLists.txt contains plugin targets" +else + echo "⚠️ CMakeLists.txt may be missing plugin targets" +fi +``` + +### 4.5 Python Bindings Check + +```bash +echo "=== Python Bindings Check ===" + +python -c " +try: + import transformer_engine + print('✅ transformer_engine imports successfully') +except ImportError as e: + print(f'❌ Import failed: {e}') + exit(1) + +# Check for plugin module +try: + from transformer_engine.pytorch import plugin + print('✅ Plugin module accessible') +except (ImportError, AttributeError) as e: + print(f'⚠️ Plugin module check: {e}') + print(' (This may be expected if plugin is loaded differently)') +" +``` + +--- + +## Level 5 — End-to-End Tests + +Only run if a GPU is available and the environment supports training. + +```bash +echo "=== Environment Check ===" +python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}'); print(f'GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"N/A\"}')" + +if python -c "import torch; exit(0 if torch.cuda.is_available() else 1)"; then + echo "=== Running E2E Test ===" + # Small-scale model training test + python -c " +import torch +import transformer_engine.pytorch as te + +# Simple forward pass with TE layers +model = te.Linear(256, 256) +x = torch.randn(2, 256, device='cuda') +with te.fp8_autocast(): + y = model(x) + loss = y.sum() + loss.backward() +print('✅ E2E forward/backward pass successful') +print(f' Output shape: {y.shape}') +" +else + echo "⚠️ No GPU available — skipping Level 5" + echo " Record this in the sync report" +fi +``` + +--- + +## Rollback Strategy + +If any level fails and cannot be fixed: + +```bash +# Find the merge commit +git log --oneline --merges -1 + +# Revert the merge (keeping the main branch's history) +git revert -m 1 + +# Push the revert +git push origin main +``` + +The `-m 1` tells git to revert to the first parent (main), effectively undoing the merge while +preserving history. diff --git a/skills/te-fl-upstream-sync/scripts/generate_sync_report.sh b/skills/te-fl-upstream-sync/scripts/generate_sync_report.sh new file mode 100644 index 0000000000..47dc71d6da --- /dev/null +++ b/skills/te-fl-upstream-sync/scripts/generate_sync_report.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +# generate_sync_report.sh — Generate a markdown sync report for TransformerEngine-FL +# Usage: bash scripts/generate_sync_report.sh [upstream_branch] [output_file] + +set -euo pipefail + +UPSTREAM_BRANCH="${1:-release_v2.14}" +OUTPUT="${2:-SYNC_REPORT.md}" + +echo "Generating sync report..." + +# Gather data +MERGE_COMMIT=$(git log --oneline --merges -1 --format="%H" 2>/dev/null || echo "N/A") +MERGE_COMMIT_SHORT=$(git log --oneline --merges -1 --format="%h" 2>/dev/null || echo "N/A") +UPSTREAM_SHA=$(git rev-parse "upstream/${UPSTREAM_BRANCH}" 2>/dev/null || echo "N/A") +UPSTREAM_SHA_SHORT=$(git rev-parse --short "upstream/${UPSTREAM_BRANCH}" 2>/dev/null || echo "N/A") +CURRENT_BRANCH=$(git branch --show-current 2>/dev/null || echo "N/A") +SYNC_DATE=$(date +"%Y-%m-%d %H:%M:%S %Z") +MAIN_HEAD=$(git rev-parse --short HEAD 2>/dev/null || echo "N/A") + +# Plugin directory status +PLUGIN_STATUS="❌ Missing" +if [ -d "transformer_engine/common/plugin" ]; then + PCOUNT=$(find transformer_engine/common/plugin -type f -name "*.h" | wc -l | tr -d ' ') + PLUGIN_STATUS="✅ Present ($PCOUNT header files)" +fi + +# CUDA patches status +CUDA_STATUS="❌ Missing" +if [ -d "transformer_engine/common/cuda_patches" ]; then + CCOUNT=$(find transformer_engine/common/cuda_patches -type f | wc -l | tr -d ' ') + CUDA_STATUS="✅ Present ($CCOUNT files)" +fi + +# OP API check +API_STATUS="❌ Not found" +if [ -d "transformer_engine/common/plugin" ]; then + REG=$(grep -rl "register_plugin" transformer_engine/common/plugin/ 2>/dev/null | wc -l | tr -d ' ') + UNREG=$(grep -rl "unregister_plugin" transformer_engine/common/plugin/ 2>/dev/null | wc -l | tr -d ' ') + if [ "$REG" -gt 0 ] && [ "$UNREG" -gt 0 ]; then + API_STATUS="✅ register_plugin() and unregister_plugin() present" + elif [ "$REG" -gt 0 ]; then + API_STATUS="⚠️ register_plugin() found, unregister_plugin() missing" + fi +fi + +# Build config check +SETUP_STATUS="N/A" +CMAKE_STATUS="N/A" +[ -f "setup.py" ] && { grep -qi "plugin" setup.py && SETUP_STATUS="✅ Plugin targets present" || SETUP_STATUS="⚠️ No plugin references"; } +[ -f "CMakeLists.txt" ] && { grep -qi "plugin" CMakeLists.txt && CMAKE_STATUS="✅ Plugin targets present" || CMAKE_STATUS="⚠️ No plugin references"; } + +# Write report +cat > "$OUTPUT" << EOF +# TransformerEngine-FL Upstream Sync Report + +## Sync Summary + +| Field | Value | +|-------|-------| +| Date | ${SYNC_DATE} | +| Upstream | Nvidia/TransformerEngine | +| Upstream Branch | ${UPSTREAM_BRANCH} | +| Upstream Commit | \`${UPSTREAM_SHA_SHORT}\` (${UPSTREAM_SHA}) | +| Merge Commit | \`${MERGE_COMMIT_SHORT}\` (${MERGE_COMMIT}) | +| Current Branch | ${CURRENT_BRANCH} | +| HEAD | \`${MAIN_HEAD}\` | + +## Plugin System Status + +| Component | Status | +|-----------|--------| +| Plugin directory | ${PLUGIN_STATUS} | +| CUDA patches | ${CUDA_STATUS} | +| OP API interfaces | ${API_STATUS} | +| setup.py | ${SETUP_STATUS} | +| CMakeLists.txt | ${CMAKE_STATUS} | + +## CI/CD Results + +| Level | Test | Status | +|-------|------|--------| +| 1 | Compile test (\`pip install -e .\`) | ⬜ Not run | +| 2 | Unit tests (\`pytest tests/unit/\`) | ⬜ Not run | +| 3 | Integration tests (\`pytest tests/pytorch/\`) | ⬜ Not run | +| 4 | Plugin-specific tests | ⬜ Not run | +| 5 | End-to-end tests | ⬜ Not run | + +> Update the CI/CD table above as each level completes. + +## Conflict Resolution Summary + +> Fill in after conflict resolution: +> +> - P0 conflicts resolved: _count_ +> - P1 conflicts resolved: _count_ +> - P2 conflicts resolved: _count_ +> - Total files with conflicts: _count_ + +## Rollback Information + +If issues are discovered post-merge: + +\`\`\`bash +git revert -m 1 ${MERGE_COMMIT} +git push origin main +\`\`\` + +## Notes + +_Add any additional notes about the sync here._ +EOF + +echo "✅ Sync report written to: $OUTPUT" +echo " Merge commit: $MERGE_COMMIT_SHORT" +echo " Upstream SHA: $UPSTREAM_SHA_SHORT" diff --git a/skills/te-fl-upstream-sync/scripts/validate_plugin.sh b/skills/te-fl-upstream-sync/scripts/validate_plugin.sh new file mode 100644 index 0000000000..43bff4a6ab --- /dev/null +++ b/skills/te-fl-upstream-sync/scripts/validate_plugin.sh @@ -0,0 +1,142 @@ +#!/usr/bin/env bash +# validate_plugin.sh — Plugin system integrity checks for TransformerEngine-FL +# Usage: bash scripts/validate_plugin.sh [repo_root] +# Exit code: 0 = all checks pass, 1 = critical failure, 2 = warnings only + +set -euo pipefail + +REPO_ROOT="${1:-.}" +FAIL=0 +WARN=0 + +echo "============================================" +echo " TransformerEngine-FL Plugin Validation" +echo "============================================" +echo "" + +# --- Check 1: Plugin directory exists --- +echo "[1/6] Plugin directory integrity" +PLUGIN_DIR="$REPO_ROOT/transformer_engine/common/plugin" +if [ -d "$PLUGIN_DIR" ]; then + FILE_COUNT=$(find "$PLUGIN_DIR" -type f | wc -l | tr -d ' ') + echo " ✅ Plugin directory exists ($FILE_COUNT files)" + find "$PLUGIN_DIR" -type f -name "*.h" | while read -r f; do + echo " - $(basename "$f")" + done +else + echo " ❌ CRITICAL: Plugin directory missing: $PLUGIN_DIR" + FAIL=1 +fi + +# --- Check 2: CUDA patches directory --- +echo "" +echo "[2/6] CUDA patches directory" +CUDA_DIR="$REPO_ROOT/transformer_engine/common/cuda_patches" +if [ -d "$CUDA_DIR" ]; then + PATCH_COUNT=$(find "$CUDA_DIR" -type f | wc -l | tr -d ' ') + echo " ✅ CUDA patches directory exists ($PATCH_COUNT files)" +else + echo " ❌ CRITICAL: CUDA patches directory missing: $CUDA_DIR" + FAIL=1 +fi + +# --- Check 3: OP API signatures --- +echo "" +echo "[3/6] OP API interface signatures" +if [ -d "$PLUGIN_DIR" ]; then + REG=$(grep -rl "register_plugin" "$PLUGIN_DIR" 2>/dev/null | wc -l | tr -d ' ') + UNREG=$(grep -rl "unregister_plugin" "$PLUGIN_DIR" 2>/dev/null | wc -l | tr -d ' ') + + if [ "$REG" -gt 0 ]; then + echo " ✅ register_plugin() found in $REG file(s)" + grep -n "register_plugin" "$PLUGIN_DIR"/*.h 2>/dev/null | head -5 | sed 's/^/ /' + else + echo " ❌ CRITICAL: register_plugin() not found" + FAIL=1 + fi + + if [ "$UNREG" -gt 0 ]; then + echo " ✅ unregister_plugin() found in $UNREG file(s)" + grep -n "unregister_plugin" "$PLUGIN_DIR"/*.h 2>/dev/null | head -5 | sed 's/^/ /' + else + echo " ❌ CRITICAL: unregister_plugin() not found" + FAIL=1 + fi +else + echo " ⏭️ Skipped (plugin directory missing)" +fi + +# --- Check 4: CUDA patches applicability --- +echo "" +echo "[4/6] CUDA patches applicability" +if [ -d "$CUDA_DIR" ]; then + PATCH_FAIL=0 + for patch in "$CUDA_DIR"/*.patch; do + [ -f "$patch" ] || continue + BASENAME=$(basename "$patch") + if git apply --check "$patch" 2>/dev/null; then + echo " ✅ $BASENAME applies cleanly" + else + echo " ⚠️ $BASENAME fails to apply (may need rebasing)" + WARN=1 + PATCH_FAIL=1 + fi + done + if [ $PATCH_FAIL -eq 0 ] && [ "$(find "$CUDA_DIR" -name '*.patch' | wc -l | tr -d ' ')" -eq 0 ]; then + echo " ℹ️ No .patch files found (patches may use a different format)" + fi +else + echo " ⏭️ Skipped (CUDA patches directory missing)" +fi + +# --- Check 5: Build configuration --- +echo "" +echo "[5/6] Build configuration" +for BUILD_FILE in setup.py CMakeLists.txt pyproject.toml; do + FPATH="$REPO_ROOT/$BUILD_FILE" + if [ -f "$FPATH" ]; then + if grep -qi "plugin" "$FPATH"; then + echo " ✅ $BUILD_FILE references plugin targets" + else + echo " ⚠️ $BUILD_FILE exists but no plugin references found" + WARN=1 + fi + else + echo " ℹ️ $BUILD_FILE not found (may not be applicable)" + fi +done + +# --- Check 6: Python bindings --- +echo "" +echo "[6/6] Python bindings" +if python -c "import transformer_engine" 2>/dev/null; then + echo " ✅ transformer_engine imports successfully" + python -c " +try: + from transformer_engine.pytorch import plugin + print(' ✅ Plugin module accessible') +except (ImportError, AttributeError, ModuleNotFoundError) as e: + print(f' ⚠️ Plugin module: {e}') + print(' (May need pip install -e . first)') +" 2>/dev/null || echo " ⚠️ Could not check plugin module" +else + echo " ⚠️ transformer_engine not installed (run pip install -e . first)" + WARN=1 +fi + +# --- Summary --- +echo "" +echo "============================================" +if [ $FAIL -gt 0 ]; then + echo " ❌ VALIDATION FAILED — critical issues found" + echo "============================================" + exit 1 +elif [ $WARN -gt 0 ]; then + echo " ⚠️ VALIDATION PASSED WITH WARNINGS" + echo "============================================" + exit 2 +else + echo " ✅ ALL CHECKS PASSED" + echo "============================================" + exit 0 +fi