Method forward

mpu/random.py:281–329

(ctx, run_function, *args)

    """
    @staticmethod
    def forward(ctx, run_function, *args):
        ctx.run_function = run_function
        global mp_rank, mp_size, mp_group
        if mp_rank is None:
            mp_rank = get_model_parallel_rank()
            mp_size = get_model_parallel_world_size()
            mp_group = get_model_parallel_group()


        global cuda_device, transport_stream, PARTITION_ACTIVATIONS
        if cuda_device is None:
            if dist.get_rank()  == 0:
                print(f"Partition Activations {PARTITION_ACTIVATIONS} and Correctness Check {PA_CORRECTNESS_TEST}")
            
            cuda_device = torch.cuda.current_device()
            #The transport stream is used to overlap the allgather communication for the activations
            #with the computation in the backward pass
            transport_stream = torch.cuda.Stream(device=cuda_device)

        if PARTITION_ACTIVATIONS:
            inputs = [item.detach().contiguous().view(-1).narrow(0, get_partition_start(item), get_partition_size(item)).clone() for item in args[:-1]]
            inputs.append(args[-1])

        #just in case something funky is happening such as reuse of inputs
        inputs_cuda = [item.to(cuda_device) for item in args]
        
        # Copy the rng states.
        ctx.fwd_cpu_rng_state = torch.get_rng_state()
        ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
        ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()

        #ctx.save_for_backward(*args)
        with torch.no_grad():
            outputs = run_function(*inputs_cuda)

        del inputs_cuda
        
        if PARTITION_ACTIVATIONS:
            new_args = []
            for arg, inp in zip(args,inputs):                
                size= torch.tensor(arg.size())
                arg.data = inp.data
                new_args.append(arg)
                new_args.append(size)
            ctx.save_for_backward(*new_args)
        else:
            ctx.save_for_backward(*args)
        
        return outputs

    @staticmethod
    def backward(ctx, *args):