get resident set size memory for the current process

Performs a single optimization step. Args: closure (callable, optional): A closure that reevaluates the model and

Clear the GPU cache for all ranks

save model and optimizer via sharded_state_dict to save_dir

apply activation checkpointing to model returns None as model is updated directly

This method loads a module from a py file which is not in the Python path

saving model via rank0 cpu streaming and full_state_dict

This function saves the train_config and FSDP config into a train_params.yaml. This will be used by converter script in the inference folder

Initialize the process group for distributed training

Set environment flags for debugging purposes

Args: params (iterable): iterable of parameters to optimize or dicts defining parameter groups

Clean up the process group after training

create date and time for file save uniqueness example: 2022-05-07-08:31:12_PM'

load local checkpoint to rank0 cpu must be called * before * passing to FSDP

load an fsdp optimizer full_state checkpoint using scatter method this ensures only rank 0 loads the optimizer state dict and scatters to other ra

save optimizer state via full state dict