client_routine_swin.py

import time
import random
import torch
import torch.nn as nn
import torch.distributed as dist
from utils.data_utils import data_loader
from utils.logs_utils import log_to_wandb, save_checkpoint_gpt, load_checkpoint_gpt, should_checkpoint, save_gradients
from utils.net_utils import create_model, load_optimizer, is_ref_worker
from communication import synchronize_all_clients, synchronize_with_rank0_first, complete_communication_step, decentralized_communication_step, wait_event, communicate_ref_worker_with_other_clients, communicate_shuffled_mask_with_other_clients, communicate_summed_mask_with_other_clients
from utils.debug_utils import *
from utils.slurm import is_slurm_job_ending
import os
import wandb
from utils.dropout_utils import *
from swin_utils.model_utils import create_model_swin
from swin_utils.scheduler_utils import ImageNetCosineScheduler
from torch.optim import SGD, AdamW
from itertools import chain
import sys
from timm.data import Mixup
from timm.loss import SoftTargetCrossEntropy
from torch.cuda.amp import GradScaler


def get_wandb_name(args, blocks, world_size):
    project_name = args.wandb_project
    entity = args.wandb_entity
    wandb_name = f"overlap_{args.num_overlaps}_{blocks} (ws = {world_size})"
        
    return entity, project_name, wandb_name

def main_client_routine(
    i_rank,
    group_ranks,
    group_params_list,
    group_buffer_list,
    group_barrier_sgd,
    group_barrier_com,
    group_barrier_param,
    group_barrier_summed_mask,
    continue_training,
    list_train_losses,
    world_size,
    batch_size,
    gpu_id_in_node,
    lr_inner,
    momentum,
    weight_decay,
    n_epoch_if_1_worker,
    H,
    count_rounds,
    graph_topology,
    graph,
    rank_gpu,
    n_gpus,
    num_overlaps,
    num_clients_in_group,
    total_rounds,
    wandb_id,
    args,
    block_selection_matrix,
):
    # if we are the first client of each gpu, we are the one in charge of the "outside" communications,
    # so we initialize the process group used for that, and corresponding cuda events

    if i_rank == 0:
        process_group = dist.init_process_group(
            backend="nccl", rank=rank_gpu, world_size=n_gpus,
        )
        print("Initializing process group in client routine on irank=", i_rank, process_group)
        
        # init a cuda event for the each p2p com to 0 inside the group + 1
        com_events = [torch.cuda.Event(blocking=True) for k in range(len(group_ranks) + 1)] # TODO WHY +1 ? 
        print("********* Group Ranks", group_ranks)
        print("******** rank gpu", rank_gpu)
        print("******** n gpus", n_gpus)
        print("*"*20)
        
    else:
        process_group = None
        # init a single cuda event for the com of our worker
        com_events = [torch.cuda.Event(blocking=True)]
        print("----- Initializing process group in client routine on irank=", i_rank, process_group)
        print("----- Group Ranks", group_ranks)
        print("----- rank gpu", rank_gpu)
        print("----- n gpus", n_gpus)
        print("----- group_buffer list", group_buffer_list)
        print("-"*20)
    
    # create the stream used for this client
    client_stream = torch.cuda.Stream()
    filestore = dist.TCPStore(os.environ.get("MASTER_ADDR"), port=int(os.environ.get("MASTER_PORT")), is_master=False)
    # all the training events happen in this client's stream

    with torch.cuda.stream(client_stream):
        
        #### INITIALIZATION ####
        
        # gather the client's rank
        rank = group_ranks[i_rank]
        # gather the client's parameters, communication buffer, and the loss mp.Value
        params = group_params_list[i_rank]
        communication_buffer = group_buffer_list[i_rank]
        train_loss = list_train_losses[i_rank]
        params_before_sgd = None
    
        
        # initialize cuda events for the synchronizations
        client_event_sgd = torch.cuda.Event(blocking=True)
        client_event_com = torch.cuda.Event(blocking=True)
        client_event_param = torch.cuda.Event(blocking=True)
        update_param_event = torch.cuda.Event(blocking=True)
        client_event_summed_mask = torch.cuda.Event(blocking=True)

        client_event_grad_start = torch.cuda.Event(blocking=True)
        client_event_grad_end = torch.cuda.Event(blocking=True)
        
        train_loader, n_batch_per_epoch = data_loader(
            rank,
            world_size,
            train=True,
            batch_size= args.batch_size,
            dataset_name=args.dataset_name,
            model_name=args.model_name,
            non_iid_data=args.non_iid_data,
            dirichlet_param=args.dirichlet_param,
        )

        test_loader, _ = data_loader(
            rank,
            world_size,
            train=False,
            batch_size= args.batch_size,
            dataset_name=args.dataset_name,
            non_iid_data=args.non_iid_data,
            dirichlet_param=args.dirichlet_param,
        )
        
        # compute the TOTAL number (sum over all workers) of gradients performed
        nb_batches_tot = n_batch_per_epoch * n_epoch_if_1_worker
        nb_batches_per_client = int(nb_batches_tot / (world_size * H))
        client_steps_per_epoch = n_batch_per_epoch // (world_size * H)
        total_rounds.value = nb_batches_per_client # Total rounds is needed for non_ref workers only. 
        print(f"n_batch_per_epoch: {n_batch_per_epoch}")
        print(f"nb_batches_tot: {nb_batches_tot}")
        print(f"nb_batches_per_client: {nb_batches_per_client}")
    
        if args.warmup:
            nb_batches_of_warmup = n_batch_per_epoch * args.warmup
            nb_batches_of_warmup_per_client = int(nb_batches_of_warmup / (world_size * H))
            print(f"nb_batches_of_warmup: {nb_batches_of_warmup}")
            print(f"nb_batches_of_warmup_per_client: {nb_batches_of_warmup_per_client}")
        else:
            nb_batches_of_warmup = 0
            nb_batches_of_warmup_per_client = 0

        first_milestone = int(0.3 * nb_batches_tot)
        second_milestone = int(0.6 * nb_batches_tot)
        third_milestone = int(0.8 * nb_batches_tot)

        print(f"milestones: [{first_milestone}, {second_milestone}, {third_milestone}]")

        save_gradient_steps = [i for i in range(11)] + [i for i in range(50, 1001, 50)] + [i for i in range(2000, 10001, 1000)] + \
            [first_milestone - 50, first_milestone, first_milestone + 50] + \
            [second_milestone - 50, second_milestone, second_milestone + 50] + \
            [third_milestone - 50, third_milestone, third_milestone + 50]
        # save_gradient_steps = []

        # create an iterator
        train_iterator = iter(train_loader)
        # Initialize the model and the optimizer
        model, criterion = create_model_swin(args.dataset_name, gpu_id_in_node, mask_in_backward_only=args.backwards_masking)

        
        mask_vector = block_selection_matrix[rank_gpu*args.n_clients_per_gpu + i_rank]
        model.model.apply_block_mask(mask_vector, args.backwards_masking)
        if rank == 0:
            model.model.print_block_mask_status()
        i_mask = model.get_parameter_mask().to(gpu_id_in_node)
        summed_mask = i_mask.clone()
        summed_mask_communication = summed_mask.clone().fill_(0.)

        synchronize_all_clients(client_event_summed_mask, client_stream, process_group, group_barrier_summed_mask, i_rank)
        communicate_summed_mask_with_other_clients(graph, group_ranks, client_stream, process_group, i_rank, summed_mask, num_clients_in_group, com_events, summed_mask_communication)
        synchronize_all_clients(client_event_com, client_stream, process_group, group_barrier_summed_mask, i_rank)

        summed_mask.copy_(summed_mask_communication)
        wait_event(client_event_summed_mask, client_stream)
        synchronize_all_clients(client_event_summed_mask, client_stream, process_group, group_barrier_summed_mask, i_rank)
        summed_mask = summed_mask.to(gpu_id_in_node)
        summed_mask[summed_mask == 0] = 1

        if rank==0:
            print("Model created", model)
            print("Model parameters:", sum(p.numel() for p in model.parameters()))
            print("Summed mask:", summed_mask)
            print("Summed mask values:", torch.unique(summed_mask))

        # loads the client's parameters into the model
        model.set_weights(params)

        # print(f'Rank: {rank} | intial weights: {model.get_weights()}')

        optimizer, scheduler = load_optimizer(
            model,
            lr_inner,
            momentum,
            weight_decay,
            args.gamma_lr_scheduler,
            args.optimizer_name,
            num_overlaps if args.adjust_lr_with_dropout else world_size,
            args.use_linear_scaling,
            nb_batches_tot,
            n_epoch_if_1_worker,
            args.batch_size,
            args.filter_bias_and_bn,
            args.dataset_name,
            args.no_lr_scheduler,
            args.cosine_schedular,
            args.step_up,
            args.min_lr,
            args.warmup_percentage
        )


        if args.dataset_name == "ImageNet":
            scheduler = ImageNetCosineScheduler(
                optimizer = optimizer,
                n_step_tot = nb_batches_tot,
                base_lr = args.lr_inner,
                warmup_lr = 5e-7,
                min_lr = args.min_lr,
                world_size = world_size,
                batch_size = args.batch_size,
                step_up = args.step_up,
                warmup_pct = args.warmup_percentage,
                cooldown_pct = args.cooldown_percentage,
            )

        loss_scaler = GradScaler()
        
        if args.resume_checkpoint is not None:
            checkpoint = load_checkpoint_gpt(args.resume_checkpoint, rank, gpu_id_in_node)
            
            if checkpoint['scheduler_step'] is not None:
                scheduler.k_step = checkpoint['scheduler_step']
            
            # if checkpoint['optimizer_state'] is not None:
                # optimizer.load_state_dict(checkpoint['optimizer_state'])
            
            if checkpoint['mask'] is not None:
                full_mask = checkpoint['mask']
            
            if checkpoint["summed_mask"] is not None:
                summed_mask = checkpoint["summed_mask"].to(gpu_id_in_node)

            if checkpoint["loss_scaler"] is not None:
                loss_scaler.load_state_dict(checkpoint["loss_scaler"])
                
            if i_rank == 0:
                count_rounds.value = checkpoint['round']
            # if checkpoint['error_buffer'] is not None:
                # error_buffer = checkpoint["error_buffer"]
                # error_buffer = error_buffer.to(gpu_id_in_node)
                
            list_train_losses[i_rank].value = checkpoint['train_loss']
            # group_params_list[i_rank].copy_(checkpoint['model_state'])
            
            print(f"Resuming from round {checkpoint['round']} with scheduler step {checkpoint['scheduler_step']}")
        

        # init the tensorboard    
        WANDB_API_KEY = os.environ.get("WANDB_API_KEY")
        WANDB_USERNAME = os.environ.get("WANDB_USERNAME")

        entity, project_name, wandb_name = get_wandb_name(args, model.get_num_blocks(), world_size)

        if int(os.environ["SLURM_PROCID"]) == 0 and i_rank == 0:
            wandb.login(key=WANDB_API_KEY)
            
            if hasattr(args, 'resume_checkpoint') and args.resume_checkpoint and checkpoint['wandb_run_id']:
                wandb.init(
                    entity=entity,
                    project=project_name,
                    id=checkpoint['wandb_run_id'],
                    resume="must",
                    name=wandb_name,
                    config=args,
                    settings=wandb.Settings(x_label="rank_0", mode="shared", x_primary=True)
                )
            else:
                run = wandb.init(
                    entity=entity,
                    id=wandb_id,
                    project=project_name,
                    name=wandb_name,
                    config=args,
                    settings=wandb.Settings(x_label="rank_0", mode="shared", x_primary=True)
                )

        
        #### TRAINING LOOP ####
        # if count_rounds.value == 0:
        #     scheduler.step() # to start at the lr from 0 when training starts from scratch.
        
        # root_dir = args.root_dir
        # if not os.path.exists(root_dir):
        #     os.makedirs(root_dir)

        if args.dataset_name == "ImageNet":
            criterion = SoftTargetCrossEntropy()
            mixup_fn = Mixup(
                mixup_alpha=0.8,    
                cutmix_alpha=1.0,   
                cutmix_minmax=None, 
                prob=1.0,           
                switch_prob=0.5,    
                mode="batch",       
                label_smoothing=0.1,
                num_classes=1000    
            )
        else:
            mixup_fn = None        


        model.train()
        # print(f" STEP {0} Before Training model params are - {model.get_weights()}", flush=True)
        for step in range(count_rounds.value, nb_batches_per_client):
            time_sgd_start = time.time()
            # iteration count
            h = 0
            lr_sum = 0

            # # If I have 15 minutes left in my slurm sript,
            # # Save a checkpoint and kill the execution process
            # if is_slurm_job_ending(15*60):
            #     if rank == 0 and step % 1000 == 0:
            #         checkpoint_dir=root_dir
            #         save_checkpoint_gpt(
            #             root_dir=root_dir,
            #             checkpoint_dir=checkpoint_dir,
            #             model=model,
            #             worker_id=rank,
            #             model_params=group_params_list[i_rank],
            #             count_rounds=count_rounds.value,
            #             train_loss=list_train_losses[i_rank].value,
            #             mask=block_selection_matrix,
            #             summed_mask = summed_mask,
            #             scheduler_step=scheduler.k_step,
            #             loss_scaler=loss_scaler,
            #             optimizer=optimizer,
            #             num_overlaps=num_overlaps,
            #         )
            #         sys.exit(0)
            #         os.kill(os.getpid(), signal.SIGTERM)
                                

            # if step > 0 and step % 1000 == 0:
            #     if rank == 0:
            #         checkpoint_dir=root_dir
            #         save_checkpoint_gpt(
            #             root_dir=root_dir,
            #             checkpoint_dir=checkpoint_dir,
            #             model=model,
            #             worker_id=rank,
            #             model_params=group_params_list[i_rank],
            #             count_rounds=count_rounds.value,
            #             train_loss=list_train_losses[i_rank].value,
            #             mask=block_selection_matrix,
            #             summed_mask = summed_mask,
            #             scheduler_step=scheduler.k_step,
            #             loss_scaler=loss_scaler,
            #             optimizer=optimizer,
            #             num_overlaps=num_overlaps,
            #         )

            #### LOCAL SGD ####
            while h < H:
                # gather the batch of data
                dataloader_start_time = time.time()
                try:
                    inputs, labels = next(train_iterator)
                except StopIteration:
                    # When the epoch ends, start a new epoch.
                    train_iterator = iter(train_loader)
                    inputs, labels = next(train_iterator)

                dataloader_end_time = time.time()
                
                # distribution of images and labels to all GPUs
                inputs = inputs.to(gpu_id_in_node, non_blocking=True)
                labels = labels.to(gpu_id_in_node, non_blocking=True)

                if mixup_fn:
                    inputs, labels = mixup_fn(inputs, labels)

                # 1) Zero gradients
                optimizer.zero_grad()

                # 2) Forward pass + loss
                with torch.amp.autocast(device_type='cuda'):
                    logits = model(inputs)
                    loss = criterion(logits, labels)

                # 3) Populate with scaled gradients
                loss_scaler.scale(loss).backward()

                # 4) Communicate the gradients (on SCALED gradients)
                if args.communicate_gradients:
                    grads = model.get_gradients()
                    synchronize_all_clients(client_event_param, client_stream, process_group, group_barrier_param, i_rank)
                    complete_communication_step(
                        i_rank,
                        i_mask,
                        graph,
                        grads,
                        communication_buffer,
                        group_ranks,
                        process_group,
                        com_events,
                        client_stream,
                        num_clients_in_group,
                    )
                    synchronize_all_clients(client_event_com, client_stream, process_group, group_barrier_com, i_rank)
                    update_local_grads(model, summed_mask, communication_buffer, update_param_event, client_stream)
                    synchronize_all_clients(client_event_param, client_stream, process_group, group_barrier_param, i_rank)

                # 5) Unscale before any grad ops
                loss_scaler.unscale_(optimizer)

                # 6) Clip the true gradients
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)

                # 7) Update model weights
                loss_scaler.step(optimizer)
                loss_scaler.update()
                
                # retrieve the lr and adds it to the sum
                lr = optimizer.param_groups[0]["lr"]
                lr_sum += lr
                # logs the loss to tensorboard
                if int(os.environ["SLURM_PROCID"]) == 0 and i_rank == 0:
                    # print("Step: {}, h: {}, Loss: {:.4f}, LR: {:.6f}".format(step, h, loss.item(), lr), flush=True)
                    log_to_wandb(H*count_rounds.value + h, rank, loss, lr, test_loader, model, torch.nn.CrossEntropyLoss() if args.dataset_name == "ImageNet" else criterion, gpu_id_in_node, args.dataset_name, delta_loss_log=5, delta_test_log=n_batch_per_epoch//4)
                    if args.stochastic_depth:
                        model.model.apply_block_mask(mask_vector, args.backwards_masking) # because in evaluate, all blocks are used
                        # if rank == 0:
                        #     model.model.print_block_mask_status()
                
                scheduler.step()
                train_loss.value = float(loss.detach().cpu())
                # update the count
                h += 1

            if not args.communicate_gradients:
            
                time_sgd_end = time.time()  # End time for local SGD
                sgd_time = time_sgd_end - time_sgd_start

        
                #### COMMUNICATION ####
                time_comm_start = time.time()
                synchronize_all_clients(client_event_sgd, client_stream, process_group, group_barrier_sgd, i_rank)
                
                # print(f" STEP {step} After SGD model params on GPU {i_rank} are - {model.get_weights()}", flush=True)

                complete_communication_step(
                        i_rank,
                        i_mask,
                        graph,
                        params,
                        communication_buffer,
                        group_ranks,
                        process_group,
                        com_events,
                        client_stream,
                        num_clients_in_group,
                    )
                
                synchronize_all_clients(client_event_com, client_stream, process_group, group_barrier_com, i_rank)
                time_comm_end = time.time()  # End time for communication
                comm_time = time_comm_end - time_comm_start

                #### UPDATE ####
                time_update_start = time.time()
            
                update_local_params(rank, summed_mask, graph, graph_topology, params, params_before_sgd, communication_buffer, False, 0, lr_sum, False, None, update_param_event, client_stream)
                synchronize_all_clients(client_event_param, client_stream, process_group, group_barrier_param, i_rank)
                            
                time_update_end = time.time()  # End time for parameter update
                update_time = time_update_end - time_update_start
                dataloader_time = dataloader_end_time - dataloader_start_time

                total_time = sgd_time + comm_time + update_time
                        
            # update the count of coms
            if i_rank == 0:
                count_rounds.value = count_rounds.value + 1
            
        # signal to the main process that the training is finished
        if i_rank == 0:
            continue_training.value = 0

@torch.no_grad()
def update_local_grads(model, summed_mask, communication_buffer, update_param_event, client_stream):
    communication_buffer.div_(summed_mask)
    # wait_event(update_param_event, client_stream)
    model.set_gradients(communication_buffer)
    # wait_event(update_param_event, client_stream)
    communication_buffer.fill_(0.)

# @torch.no_grad()
# def update_local_params(params, summed_mask, communication_buffer, update_param_event, client_stream):
#     communication_buffer.div_(summed_mask)
#     wait_event(update_param_event, client_stream)
#     params.copy_(communication_buffer)
#     wait_event(update_param_event, client_stream)
#     communication_buffer.fill_(0.)

@torch.no_grad()            
def update_local_params(rank, summed_mask, graph, graph_topology, params, params_before_sgd, communication_buffer, average_deltas, id_beg_bn, lr_sum, use_FL_momentum, mom_var, update_param_event, client_stream):
    """
    Function that update the local params of the client after local sgd and communications.
    Basic update: a "decentralized" version of FedAvg.
    Re-init all variables to their defaults values to prepare the next round.
    """
    # this is already done in the case of the complete graph
    if graph_topology != 'complete':
        # adds our version of the params to the buffer
        communication_buffer.add_(params, alpha=graph.nodes[rank]["metropolis_weight"])
    else:
        communication_buffer.div_(summed_mask)
    wait_event(update_param_event, client_stream)
    # if a form of FedAvg is performed, there is a "outer" SGD step to take for the non-BN params
    if average_deltas:
        if use_FL_momentum:
            # keep the momentum var in memory
            mom_var.copy_(communication_buffer[:id_beg_bn]).mul_(1/lr_sum)
        # replace the model's params with the previous version before local SGD
        # and take a outer grad step, using the average stored in the communication buffer
        params[:id_beg_bn].copy_(params_before_sgd[:id_beg_bn]).add_(communication_buffer[:id_beg_bn], alpha=-1)
        # average the BN params
        params[id_beg_bn:].copy_(communication_buffer[id_beg_bn:])
    # otherwise, we simply average our parameters with our peers
    else:
        params.copy_(communication_buffer)
    wait_event(update_param_event, client_stream)
    # re-initialize the com buffer
    communication_buffer.fill_(0.)
    if average_deltas:
        # save in memory the new version of the params before next local SGD round
        params_before_sgd.copy_(params)

# @torch.no_grad()
# def update_local_grads(model, summed_mask, communication_buffer):
#     communication_buffer.div_(summed_mask)
#     model.set_gradients(communication_buffer)    
#     communication_buffer.fill_(0.)