[WIP] ci: compare pipeline-parallel Eagle3 convergence

.github/workflows/gpu-integration.yml

@@ -2,7 +2,7 @@ name: GPU Integration
 
 # GPU jobs execute PR code and therefore require an explicit WRITE-access
 # maintainer trigger. A trusted runner dispatches the selected merge commit
-# into an isolated two-GPU Slurm container allocation.
+# into isolated Slurm container allocations.
 on:
   pull_request_target:
     types: [labeled]
@@ -95,12 +95,17 @@ jobs:
         uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8
         with:
           script: |
-            await github.rest.repos.createCommitStatus({
-              owner: context.repo.owner, repo: context.repo.repo,
-              sha: '${{ steps.source.outputs.head_sha }}', state: 'pending',
-              context: 'gpu-integration/train-2gpu', description: 'Waiting for GB200 runner…',
-              target_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`,
-            });
+            for (const contextName of [
+              'gpu-integration/train-2gpu',
+              'gpu-integration/train-4gpu-pp-convergence',
+            ]) {
+              await github.rest.repos.createCommitStatus({
+                owner: context.repo.owner, repo: context.repo.repo,
+                sha: '${{ steps.source.outputs.head_sha }}', state: 'pending',
+                context: contextName, description: 'Waiting for GB200 runner…',
+                target_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`,
+              });
+            }
 
   integration:
     name: Qwen3.8 Eagle3 + DSpark + DFlash2 (GB200, 2 GPUs)
@@ -154,9 +159,65 @@ jobs:
           path: ${{ runner.temp }}/slurm-report
           if-no-files-found: warn
 
+  pp_convergence:
+    name: Eagle3 TP vs PP convergence (GB200, 4 GPUs)
+    needs: authorize
+    runs-on: ${{ vars.TORCHSPEC_GPU_DISPATCHER_LABEL }}
+    timeout-minutes: 720
+    permissions:
+      contents: read
+    env:
+      TORCHSPEC_CI_MODEL: ${{ vars.TORCHSPEC_CI_MODEL || 'Qwen/Qwen3.8-27B' }}
+      TORCHSPEC_CI_MODEL_REVISION: ${{ vars.TORCHSPEC_CI_MODEL_REVISION || '1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0' }}
+      TORCHSPEC_CI_SOURCE_SHA: ${{ needs.authorize.outputs.merge_sha }}
+      TORCHSPEC_SLURM_ROOT: ${{ secrets.TORCHSPEC_SLURM_ROOT }}
+      TORCHSPEC_CI_IMAGE: ${{ secrets.TORCHSPEC_CI_IMAGE }}
+      TORCHSPEC_CI_MODEL_CACHE_HOST: ${{ secrets.TORCHSPEC_CI_MODEL_CACHE_HOST }}
+      TORCHSPEC_CI_GPU_COUNT: 4
+      TORCHSPEC_CI_RUN_SCRIPT: tools/ci/run_2gpu_training_smoke.sh
+      TORCHSPEC_CI_MODE: pp_convergence
+      TORCHSPEC_CI_PP_MAX_REL_LOSS_DIFF: 0.10
+    steps:
+      - name: Checkout trusted dispatcher
+        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09  # v5
+        with:
+          ref: ${{ github.workflow_sha }}
+          path: trusted
+          persist-credentials: false
+
+      - name: Checkout PR merge without credentials
+        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09  # v5
+        with:
+          ref: ${{ needs.authorize.outputs.merge_sha }}
+          path: source
+          persist-credentials: false
+
+      - name: Submit and wait for Slurm PP convergence test
+        run: >-
+          trusted/tools/ci/dispatch_slurm.sh
+          "$GITHUB_WORKSPACE/source"
+          "$RUNNER_TEMP/slurm-report"
+
+      - name: Publish Slurm summary
+        if: always()
+        run: |
+          if [[ -f "$RUNNER_TEMP/slurm-report/summary.md" ]]; then
+            cat "$RUNNER_TEMP/slurm-report/summary.md" >> "$GITHUB_STEP_SUMMARY"
+          else
+            echo "Slurm dispatcher exited before producing a report." >> "$GITHUB_STEP_SUMMARY"
+          fi
+
+      - name: Upload PP convergence artifacts
+        if: always()
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4
+        with:
+          name: gpu-training-4gpu-pp-${{ needs.authorize.outputs.head_sha }}
+          path: ${{ runner.temp }}/slurm-report
+          if-no-files-found: warn
+
   report:
     name: Publish GPU result
-    needs: [authorize, integration]
+    needs: [authorize, integration, pp_convergence]
     if: always() && needs.authorize.result == 'success'
     runs-on: ubuntu-latest
     permissions:
@@ -178,13 +239,27 @@ jobs:
         uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8
         with:
           script: |
-            const ok = '${{ needs.integration.result }}' === 'success';
-            await github.rest.repos.createCommitStatus({
-              owner: context.repo.owner, repo: context.repo.repo,
-              sha: '${{ needs.authorize.outputs.head_sha }}',
-              state: ok ? 'success' : 'failure',
-              context: 'gpu-integration/train-2gpu',
-              description: ok ? 'Passed on GB200' : 'Failed on GB200',
-              target_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`,
-            });
-            if (!ok) core.setFailed('GPU integration failed');
+            const results = [
+              {
+                result: '${{ needs.integration.result }}',
+                context: 'gpu-integration/train-2gpu',
+                description: 'Passed on GB200',
+              },
+              {
+                result: '${{ needs.pp_convergence.result }}',
+                context: 'gpu-integration/train-4gpu-pp-convergence',
+                description: 'Passed on GB200 with TP/PP convergence match',
+              },
+            ];
+            for (const item of results) {
+              const ok = item.result === 'success';
+              await github.rest.repos.createCommitStatus({
+                owner: context.repo.owner, repo: context.repo.repo,
+                sha: '${{ needs.authorize.outputs.head_sha }}',
+                state: ok ? 'success' : 'failure',
+                context: item.context,
+                description: ok ? item.description : 'Failed on GB200',
+                target_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`,
+              });
+            }
+            if (results.some(item => item.result !== 'success')) core.setFailed('GPU integration failed');

.github/workflows/gpu-nightly.yml

@@ -107,3 +107,58 @@ jobs:
           name: nightly-gpu-${{ github.sha }}-${{ github.run_attempt }}
           path: ${{ runner.temp }}/slurm-report
           if-no-files-found: warn
+
+  pp_convergence:
+    name: Nightly Eagle3 TP vs PP convergence (GB200, 4 GPUs)
+    needs: changes
+    if: needs.changes.outputs.should_run == 'true'
+    runs-on: ${{ vars.TORCHSPEC_GPU_DISPATCHER_LABEL }}
+    timeout-minutes: 720
+    env:
+      TORCHSPEC_CI_MODEL: Qwen/Qwen3.8-27B
+      TORCHSPEC_CI_MODEL_REVISION: 1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0
+      TORCHSPEC_CI_SOURCE_SHA: ${{ github.sha }}
+      TORCHSPEC_SLURM_ROOT: ${{ secrets.TORCHSPEC_SLURM_ROOT }}
+      TORCHSPEC_CI_IMAGE: ${{ secrets.TORCHSPEC_CI_IMAGE }}
+      TORCHSPEC_CI_MODEL_CACHE_HOST: ${{ secrets.TORCHSPEC_CI_MODEL_CACHE_HOST }}
+      TORCHSPEC_CI_GPU_COUNT: 4
+      TORCHSPEC_CI_RUN_SCRIPT: tools/ci/run_2gpu_training_smoke.sh
+      TORCHSPEC_CI_MODE: pp_convergence
+      TORCHSPEC_CI_PP_MAX_REL_LOSS_DIFF: 0.10
+    steps:
+      - name: Checkout trusted dispatcher
+        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09  # v5
+        with:
+          ref: ${{ github.workflow_sha }}
+          path: trusted
+          persist-credentials: false
+
+      - name: Checkout nightly source without credentials
+        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09  # v5
+        with:
+          ref: ${{ github.sha }}
+          path: source
+          persist-credentials: false
+
+      - name: Submit and wait for Slurm PP convergence test
+        run: >-
+          trusted/tools/ci/dispatch_slurm.sh
+          "$GITHUB_WORKSPACE/source"
+          "$RUNNER_TEMP/slurm-report"
+
+      - name: Publish Slurm summary
+        if: always()
+        run: |
+          if [[ -f "$RUNNER_TEMP/slurm-report/summary.md" ]]; then
+            cat "$RUNNER_TEMP/slurm-report/summary.md" >> "$GITHUB_STEP_SUMMARY"
+          else
+            echo "Slurm dispatcher exited before producing a report." >> "$GITHUB_STEP_SUMMARY"
+          fi
+
+      - name: Upload PP convergence artifacts
+        if: always()
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4
+        with:
+          name: nightly-gpu-4gpu-pp-${{ github.sha }}-${{ github.run_attempt }}
+          path: ${{ runner.temp }}/slurm-report
+          if-no-files-found: warn

configs/ci/vllm_qwen3_8_27b_eagle3_pp_convergence.yaml

@@ -0,0 +1,80 @@
+# Four-GPU CI lane comparing the same Eagle3 training job with vLLM TP and PP.
+#
+# Two GPUs run training and two GPUs run one vLLM target engine. The launcher
+# runs this config twice: once with pp_size=1 (TP2 reference) and once with
+# pp_size=2 (TP1 x PP2). Both runs use the same seed, fixture, and optimizer
+# step count; the launcher compares every logged TRAIN_STEP loss.
+model:
+  target_model_path: Qwen/Qwen3.8-27B
+  target_model_backend: vllm
+  draft_model_config: configs/ci/qwen3_8_27b_eagle3_draft.json
+  embedding_key: model.language_model.embed_tokens.weight
+  lm_head_key: lm_head.weight
+  norm_key: model.language_model.norm.weight
+  trust_remote_code: true
+
+dataset:
+  train_data_path: ../../examples/data/sample_conversations.jsonl
+  chat_template: qwen
+  prompt_key: conversations
+  shuffle_dataset: false
+  length_group_size: 1
+
+training:
+  attention_backend: flex_attention
+  micro_batch_size: 1
+  draft_accumulation_steps: 1
+  learning_rate: 1e-4
+  lr_total_steps: 4
+  max_concurrent_batches: 1
+  max_grad_norm: 0.5
+  max_seq_length: 16384
+  num_epochs: 1
+  num_train_steps: 4
+  prefetch_depth: 1
+  save_interval: 0
+  save_per_epoch: false
+  seed: 42
+  training_num_gpus_per_node: 2
+  training_num_nodes: 1
+  ttt_length: 2
+  warmup_ratio: 0.0
+
+inference:
+  inference_engine_type: vllm
+  store_last_hidden_states: true
+  inference_num_gpus: 2
+  inference_num_gpus_per_engine: 2
+  inference_num_gpus_per_node: 4
+  inference_batch_size: 1
+  inference_buffer_threshold: 2
+  max_sample_pool_size: 4
+  # Qwen3.8 is hybrid GDN/full-attention. Capture only full-attention layers.
+  aux_hidden_states_layers: [3, 31, 59]
+  vllm:
+    tp_size: 2
+    pp_size: 1
+    mem_fraction_static: 0.55
+    extra_args:
+      enforce_eager: true
+      language_model_only: true
+      max_model_len: 16384
+      max_num_batched_tokens: 4096
+
+mooncake:
+  master_server_address: null
+  metadata_server: null
+  protocol: tcp
+  global_segment_size: 4GB
+  local_buffer_size: 1GB
+
+output_dir: ""
+cache_dir: ../../.ci-cache/qwen3-8-27b-eagle3-pp-convergence
+model_download_dir: null
+
+logging:
+  report_to: none
+
+debug:
+  enable_perf_metrics: true
+  save_debug_train_data: null

patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/vllm_pp_hidden_states.patch

@@ -112,7 +112,50 @@ diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qw
 +            return intermediate_tensors
 
          hidden_states, _ = self.norm(hidden_states, residual)
+
+diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py
+--- a/vllm/model_executor/models/qwen3_next.py
++++ b/vllm/model_executor/models/qwen3_next.py
+@@ -653,7 +653,11 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
+         positions: torch.Tensor,
+         intermediate_tensors: IntermediateTensors | None = None,
+         inputs_embeds: torch.Tensor | None = None,
+-    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
++    ) -> (
++        torch.Tensor
++        | IntermediateTensors
++        | tuple[torch.Tensor | IntermediateTensors, list[torch.Tensor]]
++    ):
+         if get_pp_group().is_first_rank:
+             if inputs_embeds is not None:
+                 hidden_states = inputs_embeds
+@@ -666,5 +670,11 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
+         full_num_tokens = positions.shape[-1]
+-        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
++        # Capture ids are global layer positions.  Only PP0 owns the embedding
++        # output at id 0; later stages must not duplicate their stage input.
++        aux_hidden_states: list[torch.Tensor] = []
++        if self.start_layer == 0:
++            self._maybe_add_hidden_state(
++                aux_hidden_states, 0, hidden_states, residual
++            )
+         for layer_idx, layer in enumerate(
+             islice(self.layers, self.start_layer, self.end_layer),
+             start=self.start_layer,
+@@ -700,9 +710,12 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
+             )
 
+         if not get_pp_group().is_last_rank:
+-            return IntermediateTensors(
++            intermediate_tensors = IntermediateTensors(
+                 {"hidden_states": hidden_states, "residual": residual}
+             )
++            if self.aux_hidden_state_layers:
++                return intermediate_tensors, aux_hidden_states
++            return intermediate_tensors
+         if hidden_states.shape[0] != full_num_tokens:
+             hidden_states, residual = _all_gather_hidden_and_residual(
+                 hidden_states,
 diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py
 index 7efe92893..8f536e497 100644
 --- a/vllm/models/kimi_k3/nvidia/model.py
@@ -424,7 +467,7 @@ index d500b6720..78a7128c4 100644
              else:
                  # Rare case.
                  assert not self.is_pooling_model
-@@ -4542,11 +4568,23 @@ class GPUModelRunner(
+@@ -4542,11 +4568,21 @@ class GPUModelRunner(
          self, grammar_output: "GrammarOutput | None"
      ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
          if self.execute_model_state is None:
@@ -444,12 +487,10 @@ index d500b6720..78a7128c4 100644
 +                and not get_pp_group().is_last_rank
 +            ):
                  self._pp_receive_prev_sampled_token_ids_to_input_batch()
-+            if extract_hidden_states and get_pp_group().world_size > 1:
-+                get_pp_group().barrier()
              # In case of PP with kv transfer, we need to pass through the
              # kv_connector_output
              return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
-@@ -4567,6 +4605,40 @@ class GPUModelRunner(
+@@ -4567,6 +4603,38 @@ class GPUModelRunner(
          # Clear ephemeral state.
          self.execute_model_state = None
 
@@ -464,8 +505,6 @@ index d500b6720..78a7128c4 100644
 +                "sampling, token proposal, and decode"
 +            )
 +            self.finalize_kv_connector()
-+            if get_pp_group().world_size > 1:
-+                get_pp_group().barrier()
 +            self.eplb_step()
 +
 +            kv_connector_output = self.kv_connector_output
@@ -490,15 +529,3 @@ index d500b6720..78a7128c4 100644
          # Apply structured output bitmasks if present.
          if grammar_output is not None:
              apply_grammar_bitmask(
-@@ -4743,6 +4815,11 @@ class GPUModelRunner(
-         # draft model to also save its KV cache.
-         if spec_config is not None:
-             self.finalize_kv_connector()
-+            if (
-+                spec_config.uses_extract_hidden_states()
-+                and get_pp_group().world_size > 1
-+            ):
-+                get_pp_group().barrier()
-
-         with record_function_or_nullcontext("gpu_model_runner: eplb"):
-             self.eplb_step()