cannot pickle 'Stream' object

Trying to run llama3 70B

Hardware:
3x 4090

deepspeed deep.py

deep.py:

# ---------------------------------------
# New automatic tensor parallelism method
# ---------------------------------------
import os
import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import deepspeed
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))

# specify the model id
model_id = "meta-llama/Meta-Llama-3-70B-Instruct"
# create the model pipeline
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, llm_int8_enable_fp32_cpu_offload=False)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, torch_dtype="auto",device_map="auto", low_cpu_mem_usage = True)


pipe = transformers.pipeline(task="text-generation", model=model, device=local_rank)
# Initialize the DeepSpeed-Inference engine
pipe.model = deepspeed.init_inference(
    pipe.model,
    mp_size=world_size,
    dtype=torch.float
)
output = pipe('What is deepspeed')
print(output)

Traceback:

Traceback (most recent call last):
  File "/home/sw/bulldozer/code/deepspeed/DeepSpeedExamples/inference/huggingface/zero_inference/deep.py", line 17, in <module>
    model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, torch_dtype="auto",device_map="auto", low_cpu_mem_usage = True)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 555, in from_pretrained
    return model_class.from_pretrained(
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/site-packages/transformers/modeling_utils.py", line 2974, in from_pretrained
    modules_to_not_convert = get_keys_to_not_convert(model)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/site-packages/transformers/utils/bitsandbytes.py", line 257, in get_keys_to_not_convert
    tied_model = deepcopy(model)  # this has 0 cost since it is done inside `init_empty_weights` context manager`
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 297, in _reconstruct
    value = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 297, in _reconstruct
    value = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 297, in _reconstruct
    value = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 297, in _reconstruct
    value = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/sw/anaconda3/envs/deepseed/lib/python3.10/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle 'Stream' object

Reproduce the failure from DeepSpeedExamples/inference/huggingface/zero_inference/deep.py using the shown deepspeed deep.py command, then trace the AutoModelForCausalLM.from_pretrained path through the reported Transformers and bitsandbytes deepcopy call. No test is named; done means establishing a supported execution path for this example that no longer raises the reported Stream pickling error.