support multi-cards magnitude pruning

examples/pytorch/nlp/huggingface_models/language-modeling/pruning/magnitude/README.md

@@ -0,0 +1,90 @@
+Step-by-Step
+============
+
+# single GPU
+
+```
+bash run.sh
+```
+
+# multi GPU
+
+we use `accelerate` and `deepspeed ZeRO Stage-2` to conduct weight magnitude pruning
+
+### Accelerate DeepSpeed Plugin
+
+On your machine(s) just run:
+```
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for DeepSpeed to which you should answer no. Then answer the following questions to generate a basic DeepSpeed config. This will generate a config file that will be used automatically to properly set the default options when doing
+
+For instance,
+
+```
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: config/zero_stage2_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+with the contents of `config/zero_stage2_config.json` being:
+
+```
+{
+  "train_batch_size": 64,
+  "train_micro_batch_size_per_gpu": 8,
+  "gradient_accumulation_steps": 4,
+  "fp16": {
+    "enabled": true,
+    "min_loss_scale": 1,
+    "opt_level": "O2"
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_param": {
+      "device": "cpu"
+    },
+    "offload_optimizer": {
+      "device": "cpu"
+    },
+    "allgather_partitions": true,
+    "allgather_bucket_size": 5e8,
+    "contiguous_gradients": true
+  },
+  "optimizer": {
+    "type": "AdamW",
+    "params": {
+      "lr": "auto",
+      "torch_adam": true,
+      "adam_w_mode": true
+    }
+  },
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "warmup_min_lr": 0.0,
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "total_num_steps": "auto",
+      "warmup_type": "cosine"
+    }
+  }
+}
+```
+
+### pruning
+
+```
+bash run_ds.sh 
+```

examples/pytorch/nlp/huggingface_models/language-modeling/pruning/magnitude/config/zero_stage2_config.json

@@ -0,0 +1,40 @@
+{
+  "train_batch_size": 64,
+  "train_micro_batch_size_per_gpu": 8,
+  "gradient_accumulation_steps": 4,
+  "fp16": {
+    "enabled": true,
+    "min_loss_scale": 1,
+    "opt_level": "O2"
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_param": {
+      "device": "cpu"
+    },
+    "offload_optimizer": {
+      "device": "cpu"
+    },
+    "allgather_partitions": true,
+    "allgather_bucket_size": 5e8,
+    "contiguous_gradients": true
+  },
+  "optimizer": {
+    "type": "AdamW",
+    "params": {
+      "lr": "auto",
+      "torch_adam": true,
+      "adam_w_mode": true
+    }
+  },
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "warmup_min_lr": 0.0,
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto",
+      "total_num_steps": "auto",
+      "warmup_type": "cosine"
+    }
+  }
+}

examples/pytorch/nlp/huggingface_models/language-modeling/pruning/magnitude/requirements.txt

@@ -0,0 +1,8 @@
+accelerate
+datasets
+sentencepiece
+transformers
+torch
+tqdm
+optimum
+einops

examples/pytorch/nlp/huggingface_models/language-modeling/pruning/magnitude/run.sh

@@ -0,0 +1,14 @@
+export CUDA_VISIBLE_DEVICES=0
+python run_clm_no_trainer_pruning.py \
+    --dataset_name NeelNanda/pile-10k \
+    --model_name_or_path facebook/opt-125m \
+    --block_size 128 \
+    --per_device_train_batch_size 8 \
+    --gradient_accumulation_steps 4 \
+    --output_dir ./test-clm \
+    --do_prune \
+    --num_train_epochs 10 \
+    --pruning_type "magnitude" \
+    --target_sparsity 0.8 \
+    --pruning_pattern "4x1" \
+    --pruning_frequency 1000