Perforatedai Distributed

PerforatedAI Multi-GPU (Distributed) Setup Skill

This skill handles DataParallel and DistributedDataParallel (DDP) setup for PerforatedAI. It is automatically invoked by the main perforatedai skill when multi-GPU training is detected.

Prerequisites:

• User has already completed detection in main perforatedai skill (Step 4.1)
• Model initialization with UPA.perforate_model() has been added
• User confirmed they want to use DataParallel or DDP

DataParallel Setup

Use this section when user is using torch.nn.DataParallel.

Educate them about the additional complexity:

Tell them:

Important: DataParallel Setup with PerforatedAI

DataParallel with PerforatedAI requires a simple two-step initialization process:

1. First run: Initialize multi-GPU settings (exits after one batch)
2. Second run: Actual training with multi-GPU support

This is a one-time setup. After initialization, you just run your training normally.

Steps:

1. Add command-line argument to their script:

Find their argparse section (or add one if they don't have it) and add:

parser.add_argument('--perforate_model_parallel', action='store_true', 
                   help='Initialize PAI settings for multi-GPU (run once on single GPU)')

2. Add conditional initialization code at the location of their DataParallel line:

Find their DataParallel line:

model = torch.nn.DataParallel(model)

Replace it with:

# PAI multi-GPU setup - initialize on single GPU, then use DataParallel
if not args.perforate_model_parallel:
    # Normal training mode - load settings and wrap with DataParallel
    GPA.pai_tracker.initialize_tracker_settings()
    model = torch.nn.DataParallel(model)
# else: initialization mode - stay on single GPU, will save settings after first batch

2.5. Setup Optimizer and Scheduler:

Find where their optimizer and scheduler are currently defined in their script.

🚨 CRITICAL RULE: PRESERVE USER'S EXACT OPTIMIZER AND SCHEDULER TYPES AND ARGUMENTS

Replace their optimizer/scheduler code with the PAI pattern while keeping the EXACT SAME types and arguments.

Example: Original:

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

Correct PAI conversion (PRESERVES their choices):

GPA.pai_tracker.set_optimizer(torch.optim.Adam)  # SAME optimizer type
GPA.pai_tracker.set_scheduler(torch.optim.lr_scheduler.StepLR)  # SAME scheduler type
optimArgs = {'params': model.parameters(), 'lr': 0.001, 'weight_decay': 1e-4}  # SAME args
schedArgs = {'step_size': 30, 'gamma': 0.1}  # SAME scheduler args
optimizer, scheduler = GPA.pai_tracker.setup_optimizer(model, optimArgs, schedArgs)

4. Add initialization mode handling in training loop:

⚠️ IMPORTANT: This must be added INSIDE the batch/iteration loop, immediately after loss.backward(), so it exits after processing just ONE batch (not after a full epoch).

Find their batch loop (usually for batch in train_loader: or similar) and add this code immediately after loss.backward():

# Inside the batch loop:
for batch in train_loader:  # Their loop
    # ... their training code ...
    loss.backward()
    
    # PAI initialization mode - save settings and exit after FIRST batch
    if args.perforate_model_parallel:
        GPA.pai_tracker.save_tracker_settings()
        print(f"PAI multi-GPU settings saved to {{save_name}}/")
        print("Initialization complete. Now run without --perforate_model_parallel flag for training.")
        exit(0)
    
    optimizer.step()  # Their code continues
    # ...

Optional: Add gradient pre-initialization (if encountering gradient-related errors):

DataParallel typically doesn't require this, but if they encounter errors about undefined gradients, add this AFTER optimizer.zero_grad() and BEFORE loss.backward():

optimizer.zero_grad()

# Pre-initialize gradients for multi-GPU compatibility (optional)
for param in model.parameters():
    if param.requires_grad and param.grad is None:
        param.grad = torch.zeros_like(param)

loss.backward()
optimizer.step()

Tell them:

"I've set up your script for DataParallel. To train:

1. First run: python train.py --perforate_model_parallel (runs on single GPU, exits after ONE batch)
2. Second run: python train.py (uses DataParallel for multi-GPU training)
3. If you change any PAI configuration settings later, re-run step 1"

DistributedDataParallel (DDP) Setup

Use this section when user is using torch.nn.parallel.DistributedDataParallel.

Educate them about the additional complexity:

Tell them:

Important: DistributedDataParallel Setup with PerforatedAI

DistributedDataParallel with PerforatedAI requires a special workflow because when dendrites are added, the process needs to exit and restart. I'll create a shell script that automates this entire process:

• The script handles initialization and continuous training automatically
• When training is interrupted (e.g., dendrites added), it automatically restarts and resumes
• You just run one command and let it handle everything

Steps:

1. Add command-line arguments to their script:

Find their argparse section (or add one if they don't have it) and add:

parser.add_argument('--perforate_model_parallel', action='store_true', 
                   help='Initialize PAI settings for DDP (run once)')
parser.add_argument('--pai_load_folder', type=str, default=None,
                   help='Folder to load PAI state from (for automatic resumption)')

2. Add checkpoint loading logic BEFORE optimizer creation:

⚠️ CRITICAL: This must be added RIGHT AFTER model creation and perforate_model(), BEFORE creating the optimizer.

Find where they create their model and call perforate_model:

model = YourModel(...)
model = UPA.perforate_model(model, save_name="...", maximizing_score=...)
model = model.to(device)

Add this loading logic immediately after (still BEFORE any optimizer creation):

# Load checkpoint if resuming from dendrite restructure (DDP mode)
if args.pai_load_folder is not None:
    # Find the highest numbered switch_x.pt file
    import glob
    switch_files = glob.glob(f"{args.pai_load_folder}/switch_*.pt")
    if switch_files:
        # Extract switch numbers and find the maximum
        switch_numbers = []
        for f in switch_files:
            try:
                num = int(f.split('switch_')[1].split('.pt')[0])
                switch_numbers.append(num)
            except:
                pass
        if switch_numbers:
            max_switch = max(switch_numbers)
            model = UPA.load_system(model, args.pai_load_folder, f'switch_{max_switch}', True)
            print(f"Loaded PAI state from {args.pai_load_folder}/switch_{max_switch}.pt")
        else:
            print(f"Starting from beginning (no valid switch_x.pt found in {args.pai_load_folder})")
    else:
        print(f"Starting from beginning (no switch_x.pt files found in {args.pai_load_folder})")

# NEXT: Add optimizer setup HERE (see substep 2.5 below)
# THEN: Add DDP wrapper (see substep 3 below)

Why this order matters: When dendrites are added, the model structure changes. load_system loads the new structure. The optimizer needs to be created with the correct model parameters, so loading must happen BEFORE optimizer creation.

⚠️ CRITICAL ORDERING FOR YOUR SCRIPT:

1. Model creation + perforate_model() [already done]
2. Checkpoint loading [code block above] ← YOU JUST ADDED THIS
3. Optimizer creation [substep 2.5 below] ← ADD THIS NEXT
4. DDP wrapper [substep 3 below] ← THEN ADD THIS
5. Training loop modifications [substeps 4-5 below]

Substep 2.5: Setup Optimizer and Scheduler (ADD BETWEEN CHECKPOINT LOADING AND DDP WRAPPER)

Find where their optimizer and scheduler are currently defined in their script.

🚨 CRITICAL RULE: PRESERVE USER'S EXACT OPTIMIZER AND SCHEDULER TYPES AND ARGUMENTS

If their setup is clean (2-5 lines in one place):

Replace their optimizer/scheduler code with the PAI pattern while keeping the EXACT SAME types and arguments.

Example - User has Adam optimizer with StepLR scheduler: Original:

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

Correct PAI conversion (PRESERVES their choices):

GPA.pai_tracker.set_optimizer(torch.optim.Adam)  # SAME optimizer type
GPA.pai_tracker.set_scheduler(torch.optim.lr_scheduler.StepLR)  # SAME scheduler type
optimArgs = {'params': model.parameters(), 'lr': 0.001, 'weight_decay': 1e-4}  # SAME args
schedArgs = {'step_size': 30, 'gamma': 0.1}  # SAME scheduler args
optimizer, scheduler = GPA.pai_tracker.setup_optimizer(model, optimArgs, schedArgs)

⚠️ CRITICAL: Place this code in their script RIGHT AFTER the checkpoint loading block from substep 2, BEFORE the DDP wrapper from substep 3.

If their setup is complex (scattered across many lines or conditional):

Ask the user: "Your optimizer/scheduler setup is complex. Can you point me to where you want PAI's optimizer setup to go, and confirm which optimizer and scheduler types you want to use?"

3. Add conditional DDP wrapper at their DDP location:

Find their DDP initialization:

model = torch.nn.parallel.DistributedDataParallel(model, ...)

Replace it with:

# PAI DDP setup - initialize on single GPU, then use DDP
if not args.perforate_model_parallel:
    # Normal training mode - initialize settings and wrap with DDP
    GPA.pai_tracker.initialize_tracker_settings()
    
    # Wrap with DDP - IMPORTANT: find_unused_parameters=True is required for PerforatedAI
    model = torch.nn.parallel.DistributedDataParallel(
        model, 
        device_ids=[...],  # Their original device_ids if they had them
        find_unused_parameters=True  # Required for PerforatedAI dendrite growth
    )
# else: initialization mode - stay on single GPU, will save settings after first batch

IMPORTANT: If their original DDP call had other arguments (like device_ids, output_device, broadcast_buffers, etc.), preserve those arguments and add find_unused_parameters=True to them.

4. Add gradient pre-initialization for DDP compatibility:

⚠️ CRITICAL: This fixes a DDP crash caused by PerforatedAI's selective training.

Problem: PerforatedAI's selective training (Cascade Correlation) means some parameters intentionally don't receive gradients during backward(). DDP requires ALL parameters to have gradients during allreduce, or it crashes with: RuntimeError: Encountered gradient which is undefined, but still allreduced by DDP reducer.

Solution: Pre-initialize all parameter gradients to zeros AFTER optimizer.zero_grad() and BEFORE loss.backward().

🔴 EXACT PLACEMENT REQUIREMENT:

Find their training loop and locate this sequence:

optimizer.zero_grad()
loss.backward()
optimizer.step()

Insert the gradient pre-initialization code BETWEEN optimizer.zero_grad() and loss.backward():

optimizer.zero_grad()

# Pre-initialize gradients for DDP compatibility
# PerforatedAI's selective training means some params won't get gradients from backward()
# Initialize them to zero so DDP's allreduce doesn't encounter None
if args.distributed:
    for param in model.parameters():
        if param.requires_grad and param.grad is None:
            param.grad = torch.zeros_like(param)

loss.backward()
optimizer.step()

Why this placement:

1. optimizer.zero_grad() clears old gradients (prevents memory leak)
2. Pre-initialization ensures all params have gradient tensors (prevents DDP crash)
3. loss.backward() accumulates real gradients onto the zeros (correct behavior)
4. Parameters that don't receive gradients stay at zero (harmless for DDP allreduce)