update sparseGPT example (#1408)

Signed-off-by: Zhang, Weiwei1 <weiwei1.zhang@intel.com> Co-authored-by: kevinintel <hanwen.chang@intel.com>

Author: WeiweiZhang1

examples/huggingface/pytorch/language-modeling/pruning/run_clm_sparsegpt.py

@@ -22,7 +22,7 @@ def skip(*args, **kwargs):
 torch.nn.init.kaiming_uniform_ = skip
 torch.nn.init.uniform_ = skip
 torch.nn.init.normal_ = skip
-
+import re
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 import logging
@@ -69,7 +69,7 @@ def parse_args():
  parser.add_argument(
  "--calibration_dataset_name",
  type=str,
- default="wikitext-2-raw-v1",
+ default="NeelNanda/pile-10k", # e.g. wikitext-2-raw-v1
  help="The name of the pruning dataset to use (via the datasets library).",
  )
  parser.add_argument(
@@ -128,6 +128,12 @@ def parse_args():
  default=16,
  help="Batch size (per device) for the evaluation dataloader.",
  )
+ parser.add_argument(
+ "--calib_size",
+ type=int,
+ default=128,
+ help="sample size for the calibration dataset.",
+ )
  parser.add_argument(
  "--learning_rate",
  type=float,
@@ -268,10 +274,12 @@ def parse_args():
  parser.add_argument("--tasks", default=["lambada_openai"],
  help="Usually chosen with ['lambada_openai','hellaswag','winogrande','piqa']",
  )
- parser.add_argument("--eval_fp16", action='store_true',
- help=" fp16")
  parser.add_argument("--use_accelerate", action='store_true',
- help="Usually use to accelerate evaluation for large models")
+ help="Usually use to accelerate evaluation for large models"
+ )
+ parser.add_argument("--eval_dtype", default='fp32',
+ help="choose in bf16, fp16 and fp32"
+ )
 
  args = parser.parse_args()
 
@@ -376,34 +384,33 @@ def main():
  logger.warning("You are instantiating a new config instance from scratch.")
 
  is_llama = bool("llama" in args.model_name_or_path)
- is_t5 = bool("t5" in args.model_name_or_path)
  if args.tokenizer_name:
  tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
  elif args.model_name_or_path:
  if is_llama:
  tokenizer = transformers.LlamaTokenizer.from_pretrained(args.model_name_or_path)
  else :
- tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer, trust_remote_code=True)
+ tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, 
+ use_fast=not args.use_slow_tokenizer, trust_remote_code=True)
  else:
  raise ValueError(
  "You are instantiating a new tokenizer from scratch. This is not supported by this script."
  "You can do it from another script, save it, and load it from here, using --tokenizer_name."
  )
 
  if args.model_name_or_path:
- if is_t5:
- model = T5ForConditionalGeneration.from_pretrained(
- args.model_name_or_path,
- config=config,
- )
+ if re.search("chatglm", args.model_name_or_path.lower()):
+ model = AutoModel.from_pretrained(args.model_name_or_path, 
+ trust_remote_code=args.trust_remote_code) # .half()
  else:
  model = AutoModelForCausalLM.from_pretrained(
  args.model_name_or_path,
  from_tf=bool(".ckpt" in args.model_name_or_path),
  config=config,
  trust_remote_code=args.trust_remote_code,
- low_cpu_mem_usage=args.low_cpu_mem_usage,
+ low_cpu_mem_usage=args.low_cpu_mem_usage
  )
+ 
 
  else:
  logger.info("Training new model from scratch")
@@ -492,7 +499,7 @@ def group_texts(examples):
  train_dataset = lm_datasets["train"]
 
  # DataLoaders creation:
- train_dataset = train_dataset.shuffle(seed=42).select(range(128))
+ train_dataset = train_dataset.shuffle(seed=42).select(range(args.calib_size))
  total_batch_size = args.per_device_train_batch_size
  if local_rank != -1:
  total_batch_size *= WORLD_SIZE
@@ -543,8 +550,10 @@ def group_texts(examples):
  torch.backends.cudnn.allow_tf32 = False
  use_cache = model.config.use_cache
  model.config.use_cache = False
- 
+ import time
+ s = time.time()
  pruning = prepare_pruning(model, configs, dataloader=train_dataloader, device=device)
+ logger.info(f"cost time: {time.time() - s}")
  model.config.use_cache = use_cache
 
  if args.output_dir is not None:
@@ -555,20 +564,28 @@ def group_texts(examples):
  logger.info(f"The model has been exported to {output_dir}")
 
  if device != 'cpu':
- model = model.to(device)
+ if not args.use_accelerate:
+ model = model.to(device)
+ else:
+ model = model.cpu()
  logger.info(f"***** Evaluation in GPU mode. *****")
  else:
  logger.info(f"***** Evaluation in CPU mode. *****")
  model.eval()
 
  model_name = args.model_name_or_path
- dtype = 'float32'
- if args.eval_fp16:
- if (hasattr(model, 'config') and model.config.torch_dtype is torch.bfloat16):
- dtype = 'bfloat16'
- else:
- dtype='float16'
- model_args = f'pretrained={model_name},tokenizer={model_name},dtype={dtype},use_accelerate={args.use_accelerate}'
+ dtype = None
+ if args.eval_dtype == 'bf16':
+ model = model.to(dtype=torch.bfloat16)
+ dtype = 'bfloat16'
+ elif args.eval_dtype == 'fp16':
+ dtype = 'float16'
+ model = model.to(dtype=torch.float16)
+ else:
+ dtype = 'float32'
+ model = model.to(dtype=torch.float32)
+ 
+ model_args = f'pretrained={model_name},tokenizer={model_name},dtype={dtype},use_accelerate={args.use_accelerate},trust_remote_code={args.trust_remote_code}'
  eval_batch = args.per_device_eval_batch_size
  user_model = None if args.use_accelerate else model
  results = evaluate(

examples/huggingface/pytorch/language-modeling/pruning/scripts/run_llm_sparseGPT.sh

@@ -11,10 +11,11 @@ export CUBLAS_WORKSPACE_CONFIG=':4096:8'
 #cd intel-extension-for-transformers
 python examples/huggingface/pytorch/language-modeling/pruning/run_clm_sparsegpt.py \
  --model_name_or_path /PATH/TO/LLM/ \
- --calibration_dataset_name wikitext-2-raw-v1 \
  --do_prune \
  --device=0 \
  --output_dir=/PATH/TO/SAVE/ \
+ --eval_dtype 'bf16' \
+ --per_device_eval_batch_size 16 \
  --target_sparsity 0.5 \
  --pruning_pattern 1x1