split the tensor allocate operation of comm tensors from function prepComms to function generate_io_tensors

[TaekyungHeo]

Summary

split the tensor allocate operation of comm tensors from function prepComms to function generate_io_tensors

Test Plan

1. test use 43B megatron trace full replay with storage recycle
   slurm script:

#!/bin/bash
#SBATCH --job-name=full_replay_job
#SBATCH --nodes=4
#SBATCH --ntasks=32
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --time=01:00:00
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err


srun -A general_cs_plutus\
     --container-image /lustre/fsw/portfolios/general/users/songyant/docker/songyant+docker-images+ai_develop.sqsh \
     --no-container-mount-home \
     --container-mounts=/lustre/fsw/portfolios/general/users/songyant/workspace_local:/workspace_local \
     zsh -c "
        if [ ! -d /run/sshd ]; then
            mkdir -p /run/sshd
            chmod 0755 /run/sshd
        fi

        source ~/.zshrc
        conda activate full_et_replay

        et_replay --trace_path /workspace_local/full_et_replay/traces_mega_10_new/execution_trace --recycle_storages
     "

result:

2. test use 8B megatron trace full replay with storage recycle
   slurm script:

#!/bin/bash
#SBATCH --job-name=full_replay_job
#SBATCH --nodes=2
#SBATCH --ntasks=16
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --time=01:00:00
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err


srun -A general_cs_plutus\
     --container-image /lustre/fsw/portfolios/general/users/songyant/docker/songyant+docker-images+ai_develop.sqsh \
     --no-container-mount-home \
     --container-mounts=/lustre/fsw/portfolios/general/users/songyant/workspace_local:/workspace_local \
     zsh -c "
        if [ ! -d /run/sshd ]; then
            mkdir -p /run/sshd
            chmod 0755 /run/sshd
        fi

        source ~/.zshrc
        conda activate full_et_replay

        et_replay --trace_path /workspace_local/full_et_replay/traces_mega_8_new/execution_trace --recycle_storages
     "

result:

3. test use hf_GPT2_et.json compute only replay with storage recycle

command:

et_replay --input /workspace_local/full_et_replay/traces/hf_GPT2_et.json --compute --recycle_storages

result: