test: add comprehensive DistributedTraining integration tests

[ooples]

Summary

• Add 47 comprehensive integration tests for the DistributedTraining module
• Fix 3 bugs discovered through testing: DivideByZeroException in PipelineParallelModel, TensorParallelModel, and ArgumentException in HybridShardedModel

Test Coverage

• InMemoryCommunicationBackend: AllReduce, AllGather, Broadcast, Scatter, Barrier, Send/Receive operations
• DDPModel: Constructor, metadata, parameter gathering
• FSDPModel: Constructor, metadata, sharding
• ZeRO-1/2/3 Models: Full lifecycle tests with sharding verification
• PipelineParallelModel: Constructor, metadata, pipeline stage verification
• TensorParallelModel: Constructor, metadata, tensor parallelism verification
• HybridShardedModel: Constructor, metadata, 3D parallelism configuration

Bug Fixes

Bug	Root Cause	Solution
DivideByZeroException in PipelineParallelModel	_numStages was 0 when InitializeSharding() was called	Added OnBeforeInitializeSharding() to set fields from Config
DivideByZeroException in TensorParallelModel	_tensorParallelSize was 0 when InitializeSharding() was called	Added OnBeforeInitializeSharding() to set fields from Config
ArgumentException in HybridShardedModel	Constructor params not accessible in OnBeforeInitializeSharding()	Used static ThreadLocal to pass params before base constructor

Root Cause: C# constructor initialization order - derived class fields are uninitialized when the base constructor calls virtual methods.

Test Plan

• All 47 DistributedTraining integration tests pass
• Build succeeds on both net10.0 and net471 targets
• No regressions in existing tests

Closes #654

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Review	Updated (UTC)
aidotnet-playground-api	Ready	Preview, Comment	Jan 23, 2026 8:37pm

[coderabbitai]

Summary by CodeRabbit

• Refactor

  • Restructured distributed training initialization flow to support deferred configuration setup in hybrid sharding patterns.
  • Updated parallelism models to defer field initialization until pre-sharding phase.

• Tests

  • Added comprehensive integration test suite for distributed training, covering DDP, FSDP, ZeRO, Pipeline, Tensor, and Hybrid parallelism strategies.

• Chores

  • Removed test project dependency.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Walkthrough

This PR introduces a deferred initialization hook (OnBeforeInitializeSharding) in the distributed training module's base class, allowing derived classes to configure parameters before sharding initialization. It also adds comprehensive integration tests covering multiple parallelism strategies (DDP, FSDP, ZeRO, Pipeline, Tensor, Hybrid).

Changes

Cohort / File(s)	Summary
Initialization Hook & Mechanism src/DistributedTraining/ShardedModelBase.cs, src/DistributedTraining/HybridShardedModel.cs	Introduces protected virtual OnBeforeInitializeSharding() hook called before InitializeSharding(). Implements ThreadLocal-based PendingConfig mechanism in HybridShardedModel to defer constructor parameter passing, enabling configuration in the new hook while respecting C# base-constructor constraints. Reworks ranking/shard-layout calculations and gradient synchronization flow.
Deferred Field Initialization src/DistributedTraining/PipelineParallelModel.cs, src/DistributedTraining/TensorParallelModel.cs	Converts readonly fields (_stageId, _numStages, _tensorParallelSize) to mutable fields initialized in OnBeforeInitializeSharding() overrides, deferring setup until after base hook invocation.
Integration Test Suite tests/AiDotNet.Tests/IntegrationTests/DistributedTraining/DistributedTrainingIntegrationTests.cs	Adds 1227 lines of comprehensive integration tests covering InMemoryCommunicationBackend, ShardingConfiguration, DDP/FSDP/ZeRO/Pipeline/Tensor/Hybrid model behaviors, parameter handling, serialization, and edge cases. Includes nested MockDistributedModel helper for test support.
Project Dependency Cleanup tests/AiDotNet.Tests/AiDotNetTests.csproj	Removes ProjectReference to AiDotNet.Native.CLBlast.

Sequence Diagram(s)

sequenceDiagram
    actor DC as Derived Class<br/>(HybridShardedModel)
    participant PendingConfig as ThreadLocal<br/>PendingConfig
    participant Base as ShardedModelBase<br/>Constructor
    participant Hook as OnBeforeInitialize<br/>Sharding()
    participant InitShard as InitializeSharding()

    DC->>PendingConfig: StoreConfigAndPassThrough<br/>(pipeline, tensor, data sizes)
    PendingConfig-->>DC: Store values in ThreadLocal
    DC->>Base: Call base constructor
    Base->>Hook: Call OnBeforeInitializeSharding()
    Hook->>PendingConfig: Read deferred config
    PendingConfig-->>Hook: Return stored values
    Hook->>Hook: Initialize fields from<br/>PendingConfig values
    Base->>InitShard: Call InitializeSharding()
    InitShard->>InitShard: Use initialized fields<br/>for sharding logic
    Base-->>DC: Constructor completes
    DC->>PendingConfig: Clear ThreadLocal

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• Work on issue 309 and gather info #393: Directly related — the initialization flow changes (OnBeforeInitializeSharding hook and ThreadLocal PendingConfig mechanism) modify the same core classes (HybridShardedModel, PipelineParallelModel, TensorParallelModel, ShardedModelBase) and initialization semantics that were introduced in that PR.

Suggested labels

feature

Poem

🐰 A hop through deferred dreams,
Where constructors pass via ThreadLocal streams,
OnBeforeShard lets configs align,
Before initialization's design,
Integration tests ensure all's fine! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 11.39% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly and concisely describes the main objective of the PR: adding comprehensive integration tests for DistributedTraining module.
Description check	✅ Passed	The description comprehensively explains the changes: 47 integration tests, 3 bug fixes with root causes and solutions, test plan results, and closing issue #654.
Linked Issues check	✅ Passed	The PR addresses all objectives from #654: adds comprehensive integration tests for DistributedTraining module covering multi-GPU coordination, gradient synchronization, data/model parallelism, and communication backends, with 47 tests achieving substantial coverage.
Out of Scope Changes check	✅ Passed	All changes are within scope: test file additions, bug fixes in source files directly related to issues found during testing, and one minor dependency removal in test project.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch test/distributed-training-integration-tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

In general, to fix a “virtual call in constructor” problem, you move the virtual call out of the constructor and into a separate initialization step that is invoked after the object is fully constructed, or you replace the virtual hook with a non‑virtual callback (e.g., passing a delegate or using a factory). For this class, we should ensure that OnBeforeInitializeSharding() is not invoked from the base constructor while still preserving the intended initialization sequence: (1) base state, (2) subclass custom pre‑sharding logic, (3) sharding.

The least intrusive and safest approach, given we can only edit this file, is:

1. Introduce a new non‑virtual instance method (e.g., Initialize()) that performs the current constructor’s “late” work: it calls OnBeforeInitializeSharding() and InitializeSharding(). This method can remain protected so that external callers don’t depend on it unless they choose to.
2. Remove the direct calls to OnBeforeInitializeSharding() and InitializeSharding() from the constructor so that the base constructor only initializes base fields and the communication backend.
3. Document (via comments) that derived classes are expected to invoke Initialize() after construction, or that some external factory/creator will do so. Since we cannot edit other files, we must avoid assuming or changing external call sites; instead, we preserve behavior by providing the new method but ensure the constructor no longer makes the problematic virtual call.

Within src/DistributedTraining/ShardedModelBase.cs, this means:

• Keep the constructor’s field initialization (lines 121–135) as‑is.
• Remove lines 137–140 that call OnBeforeInitializeSharding(); InitializeSharding();.
• Add a new protected void Initialize() method (or similarly named) right after the constructor that performs those two calls in the same order. This method can safely be called by derived classes at the end of their constructors when they are fully initialized.

No new imports or external dependencies are required; the change is purely structural within this class.