How to load a tensor+pipeline parallel checkpoint for inference tasks?

Following the training script here as a template:

https://github.com/bigscience-workshop/bigscience/blob/master/train/tr1-13B-base/tr1-13B-round1.slurm

I've trained some models using 2-way tensor parallelism and 4-way pipeline parallelism, which produces a number of checkpoints in directories like "global_step26000".

I'm now trying to use one of those trained checkpoints to do inference. In particular, I'm trying to work with modified versions of these scripts where I can provide new strings that are tokenized and processed on the fly:

https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/examples/generate_text.sh
https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/tools/generate_samples_gpt.py

I've tried a number of different approaches, but no luck so far. I also can't seem to find any instructions written up on that.

Anyone know the steps to load one of those checkpoints for inference on new data?