Aurora > Data Science > Frameworks > PyTorch, oneCCL pages need updates

[felker]

Description

@khossain4337 made several updates to the PyTorch page in October in 29d9187, but not every change made it to the oneCCL page in 88afc24

Both pages could use some review and improvements, since they haven't had many editors this year. We should include more folks beyond just @FilippoSimini @khossain4337 and @rickybalin

Possible updates needed, especially to https://docs.alcf.anl.gov/aurora/data-science/frameworks/oneCCL/

• Remove all mentions of Horovod? Or at least the PyTorch + Horovod example?
• Remove IPEX mentions?
• Check that the oneCCL env vars are set to the latest recommendations
• Deduplicate example scripts between the two pages. Create .py file stored in that directory and included in each via https://squidfunk.github.io/mkdocs-material/setup/extensions/python-markdown-extensions/#snippets

This came up because @brianhol42 adapted the OneCCL > DDP example to an FSDP2 example, and he was getting hangs between two nodes. Unclear if there is a bug in the stack, his code, or the documentation was missing some vital environment variable setting etc.

See argonne-lcf/test_frameworks#7

Decisions, questions, uncertainties, and dependencies

No response

Priority Level

Medium (Normal priority)

Label Confirmation

• I have selected at least one appropriate label

[felker]

@brianhol42's FSDP2 example:

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard
import os

from mpi4py import MPI

def setup(rank, local_rank, world_size):
    print(f"Hello from {rank} ({local_rank}) of {world_size}", flush=True)
    MPI.COMM_WORLD.Barrier()
    dist.init_process_group("xccl", init_method='env://', rank=rank, world_size=world_size)
    print(f"init_finished {rank}",flush=True)
    torch.xpu.set_device(local_rank)

    dist.barrier()

class SmallModel(nn.Module):
    def __init__(self):
        super(SmallModel, self).__init__()
        self.layer1 = nn.Linear(128,16)
        self.layer2 = nn.Linear(16,1)

    def forward(self, x):
        return self.layer2(self.layer1(x))

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.layer1 = nn.Linear(128,1024)
        self.layer2 = nn.Linear(1024,1024)
        self.layer3 = nn.Linear(1024,128)
        self.layer4 = nn.Linear(128,1)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

def main():
    #print(os.environ)
    local_rank=int(os.environ["PALS_LOCAL_RANKID"])
    rank = int(os.environ["PMIX_RANK"])
    #world_size = int(os.environ["PALS_LOCAL_SIZE"])
    world_size = int(os.environ["WORLD_SIZE"]) #SET THIS IN LAUNCH SCRIPT
    setup(rank,local_rank,world_size)

    print(os.environ["MASTER_ADDR"],flush=True)
    print(os.environ["MASTER_PORT"],flush=True)

    #model = SimpleModel().xpu(local_rank)
    model = SmallModel().xpu(local_rank)

    #if rank==0:
    #    for name, param in model.named_parameters():
    #        print(f"Name: {name}, Shape: {param.shape}, Shard Info: {param.placements}, Device: {param.device}")

    try:
        print('')
        fully_shard(model.layer1)
        fully_shard(model.layer2)
        #fully_shard(model.layer3)
        #fully_shard(model.layer4)
    except Exception as err:
        print(f"{err=}")
    fully_shard(model)

    print(f"model sharded {rank}")

    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    criterion = nn.MSELoss()

    for epoch in range(5):
        data = torch.randn(2048, 128).xpu(local_rank)
        labels = torch.randn(2048,1).xpu(local_rank)

        outputs = model(data)
        loss = criterion(outputs, labels)

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

        if rank==0:
            print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}",flush=True)

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

Executed with:
WORLD_SIZE=2 mpiexec -np 2 -ppn 1 python3 my_fsdp_model.py

Sourcing this file with env var settings first:

#!/bin/bash -l

module load frameworks

export CCL_PROCESS_LAUNCHER=pmix
export CCL_ATL_TRANSPORT=mpi
export CCL_ALLREDUCE_SCALEOUT="direct:0-1048576;rabenseifner:1048577-max"  # currently best allreduce algorithm at large scale
export CCL_BCAST=double_tree # currently best bcast algorithm at large scale

export CCL_KVS_MODE=mpi
export CCL_CONFIGURATION_PATH=""
export CCL_CONFIGURATION=cpu_gpu_dpcpp
export CCL_KVS_CONNECTION_TIMEOUT=600

export CCL_ZE_CACHE_OPEN_IPC_HANDLES_THRESHOLD=1024
export CCL_KVS_USE_MPI_RANKS=1

export MPI_PROVIDER=$FI_PROVIDER
unset MPIR_CVAR_CH4_POSIX_COLL_SELECTION_TUNING_JSON_FILE
unset MPIR_CVAR_CH4_COLL_SELECTION_TUNING_JSON_FILE
unset MPIR_CVAR_COLL_SELECTION_TUNING_JSON_FILE

export MASTER_ADDR=$(head -n 1 ${PBS_NODEFILE})
export MASTER_PORT=29500

[khossain4337]

I have been able to reproduce the hang. It seems simply adding extra environmental variables that worked for similar examples with DDP is not working for this case. Please allow me a little time to investigate further.
@felker @brianhol42

[brianhol42]

@felker @khossain4337 I've synced with Khalid and we have a viable workaround for the multi-node fsdp issue, which is to move to oneDPL 2021.17. I'm trying to figure out the best way to submit a JIRA ticket for intel based on the current installed oneCCL version.

Thank you to @khossain4337 for debugging this!

[khossain4337]

Many thanks, @brianhol42!

Just fixing a typo: oneDPL/2021.17 --> oneCCL/2021.17

[brianhol42]

Many thanks, @brianhol42!

Just fixing a typo: oneDPL/2021.17 --> oneCCL/2021.17

Thanks! Fixed the typo