[BUG] buffer nbytes != file bytes when using DeepSpeed and NVME offloading

[KMFODA]

Describe the bug
Running DeepSpeed with NVME offloading using the following config file results in the following error even though as per this issue the buffer size is set as a multiple of the bucket size:

../../workspace/zero_stage_3/fp16params/rank3/1_param.tensor.swp: buffer nbytes != file bytes 16777216 != 8388608
python3: /usr/local/lib/python3.6/dist-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp:223: int deepspeed_aio_handle_t::pread(const at::Tensor&, const char*, bool, bool): Assertion `static_cast<long long int>(buffer.nbytes()) == num_file_bytes' failed.
../../workspace/zero_stage_3/fp16params/rank1/1_param.tensor.swp: buffer nbytes != file bytes 16777216 != 8388608
python3: /usr/local/lib/python3.6/dist-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp:223: int deepspeed_aio_handle_t::pread(const at::Tensor&, const char*, bool, bool): Assertion `static_cast<long long int>(buffer.nbytes()) == num_file_bytes' failed.
../../workspace/zero_stage_3/fp16params/rank2/1_param.tensor.swp../../workspace/zero_stage_3/fp16params/rank0/1_param.tensor.swp: buffer nbytes != file bytes : buffer nbytes != file bytes 1677721616777216 !=  != 83886088388608

python3: /usr/local/lib/python3.6/dist-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp:223: int deepspeed_aio_handle_t::pread(const at::Tensor&, const char*, bool, bool): Assertion `static_cast<long long int>(buffer.nbytes()) == num_file_bytes' failed.
python3: /usr/local/lib/python3.6/dist-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp:223: int deepspeed_aio_handle_t::pread(const at::Tensor&, const char*, bool, bool): Assertion `static_cast<long long int>(buffer.nbytes()) == num_file_bytes' failed.
[2022-06-09 21:10:12,429] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 3616
[2022-06-09 21:10:12,430] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 3617
[2022-06-09 21:10:12,430] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 3618
[2022-06-09 21:10:12,430] [INFO] [launch.py:178:sigkill_handler] Killing subprocess 3619
[2022-06-09 21:10:12,430] [ERROR] [launch.py:184:sigkill_handler] ['/usr/bin/python3', '-u', 'run_summarization.py', '--local_rank=3', '--deepspeed', 'ds_config_zero3_6.json', '--model_name_or_path', 'allenai/led-large-16384', '--per_device_train_batch_size', '2', '--output_dir', 'output_dir', '--overwrite_output_dir', '--do_train', '--predict_with_generate', '--report_to', 'wandb', '--load_best_model_at_end', 'True', '--greater_is_better', 'True', '--evaluation_strategy', 'steps', '--metric_for_best_model', 'rouge_average', '--pad_to_max_length', 'True', '--max_source_length', '8192', '--generation_max_length', '512', '--save_steps', '1200', '--eval_steps', '400', '--logging_steps', '400', '--dataset_name', 'kaizan/amisum_v1', '--learning_rate', '0.00005', '--num_train_epochs', '10', '--weight_decay', '0.5'] exits with return code = -6

To Reproduce
Steps to reproduce the behavior:

1. git clone https://github.com/huggingface/transformers.git
2. huggingface-cli login3.
3. sed -i 's/load_optimizer_states=True/load_optimizer_states=False/g' ../transformers/src/transformers/trainer.py
4. sed -i 's/load_lr_scheduler_states=True/load_lr_scheduler_states=False/g' ../transformers/src/transformers/trainer.py
5. create a json file called ds_config_zero.json with the following ds variables assigned:

{

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },


    "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {
                "device": "nvme",
                "nvme_path": "../../workspace",
                "pin_memory": true,
                "buffer_count": 4,
                "fast_init": false
            },
            "offload_param": {
                "device": "nvme",
                "nvme_path": "../../workspace",
                "pin_memory": true,
                "buffer_count": 5,
                "buffer_size": 99614720,
                "max_in_cpu": 998244352
            },
            "aios": {
                "block_size": 262144,
                "queue_depth": 32,
                "thread_count": 1,
                "single_submit": false,
                "overlap_events": true
            },
            "overlap_comm": true,
            "contiguous_gradients": true,
            "sub_group_size": 1048576,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": "auto",
            "stage3_max_live_parameters": 1e9,
            "stage3_max_reuse_distance": 1e9,
            "stage3_gather_16bit_weights_on_model_save": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

6. run the following code:

deepspeed transformers/examples/pytorch/summarization/run_summarization.py
   --deepspeed ds_config_zero3.json \
    --model_name_or_path allenai/led-large-16384 \
    --per_device_train_batch_size 2 \
    --output_dir output_dir \
    --overwrite_output_dir \
    --do_train \
    --predict_with_generate \
    --report_to wandb \
    --load_best_model_at_end True \
    --greater_is_better True \
    --evaluation_strategy steps \
    --metric_for_best_model rouge_average \
    --pad_to_max_length True \
    --max_source_length 1024 \
    --generation_max_length 512 \
    --save_steps 1200 \
    --eval_steps 400 \
    --logging_steps 400 \
    --dataset_name kaizan/amisum_v1 \
    --learning_rate 0.00005 \
    --num_train_epochs 10 \
    --weight_decay 0.5

Expected behavior
Expected to download the model, parallelise across 4 GPUs and then start training whilst offloading parameters to NVME storage

ds_report output

[2022-06-09 22:18:31,012] [WARNING] [partition_parameters.py:54:<module>] unable to find torch.distributed._all_gather_base. will fall back to torch.distributed.all_gather which will result in suboptimal performance. please consider upgrading your pytorch installation.
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/usr/local/lib/python3.6/dist-packages/torch']
torch version .................... 1.8.0
torch cuda version ............... 10.2
torch hip version ................ None
nvcc version ..................... 10.2
deepspeed install path ........... ['/usr/local/lib/python3.6/dist-packages/deepspeed']
deepspeed info ................... 0.6.1, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.8, cuda 10.2, hip 0.0

System info (please complete the following information):

• OS = Linux
• GPU count = 4 TeslaV100S
• Python = Python 3.6.9
• Any other relevant info about your setup

Launcher context
deepspeed launcher

Docker context
N/A

Additional context
N/A

[tjruwase]

@KMFODA, can you repro with #2016?

[KMFODA]

This seems to have been fixed. Will close for now and re-open if I it ever re-occurs. Thanks for the help and prompt responses!

[niebowen666]

@KMFODA Sorry, I have suffered the same situation. Hoe did you solve this issue? Thanks!