Inference Using DeepSpeed

The purpose of this document is to guide Data Scientists to run inference on pre-trained PyTorch models using a DeepSpeed interface integrated with Gaudi.

DeepSpeed Validated Configurations

The following DeepSpeed configurations have been validated to be fully functioning with Gaudi:

Configuration	Description
Inference on multi-node	Run the model on two or more independent servers.
Tensor Parallel	Splits the model tensors into chunks so that each tensor resides on its designated HPU.
BF16 precision	Reduces model memory consumption and improves performance.
Weights sharding using meta device	Avoids loading all weights to RAM. Instead, it loads the weights of each device to RAM separately and copies to the device to be able to fit into host memory.
CUDA graph (with HPU Graph implementation)	DeepSpeed provides a flag for capturing the CUDA-Graph of the inference ops. Using the graph replay, the graphs run faster. To enable this flag, see Using HPU Graphs.

Note

DeepSpeed’s multi-node training uses pdsh for invoking the processes on remote hosts. Make sure it is installed on your machine before using it.

Installing DeepSpeed Library

The HabanaAI GitHub has a fork of the DeepSpeed library that includes changes to add support for SynapseAI. To use DeepSpeed with Gaudi, you must install Habana’s fork for DeepSpeed directly from the DeepSpeed fork repository located in HabanaAI GitHub:

pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.11.0

DeepSpeed inference was tested on this fork which is based on DeepSpeed v0.9.4.

Integrating DeepSpeed with Gaudi

Follow DeepSpeed’s instructions at https://www.deepspeed.ai/tutorials/inference-tutorial/ with the following modifications:

• Pass an additional parameter args.device = ‘hpu’ to deepspeed.init_inference(). For example:

  world_size = int(os.getenv('WORLD_SIZE', '0'))
  parser = argparse.ArgumentParser()
  args = parser.parse_args(args='')
  args.device = "hpu"
  model = deepspeed.init_inference(model,
                                   mp_size=world_size,
                                   dtype=dtype,
                                   args=args,
                                   **kwargs)
  

• Run the DeepSpeed runner:

  deepspeed --num_gpus <amount of gaudi cards> model.py <model args>
  

Multi-node is also supported. Refer to DeepSpeed documentation for more details.

Using HPU Graphs

Habana’s fork for DeepSpeed supports HPU graph natively. It is enabled by setting enable_cuda_graph=True when the device is set to ‘hpu’.

To enable HPU Graphs:

• Set the following environment variable: export PT_HPU_ENABLE_LAZY_COLLECTIVES=true

• Pass the following parameter to deepspeed initialization: deepspeed.init_inference(enable_cuda_graph=True,...)

For more details on HPU Graphs, refer to Run Inference Using HPU Graphs.

Reference Model

For a full example of a DeepSpeed inference model with Gaudi, refer to BLOOM Model Reference.