[REQUEST] Is it possible to run zero inference with cpu offloading on two separate computers?

We have two computers, each of them with an NVIDIA 4090, after following the docs, we were able to finally get the zero inference running with CPU parameter offloading. However, when configured through accelerate to run on two computers for distributed inference we are being unable to make it work. Our work is essentially a compendium of different tutorials, what we see is that both graphic cards take a relatively small gpu usage like 2 gigs when they are 24 gb each, and there is no performance gain when running in distributed mode, i.e. through accelerate launch --config_file ./config-geforfuerte02.yaml ../demo.py --output_dir /tmp I know we are missing something, but at this point i am a mess with different terms that i am not familiar with yet

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import torch.distributed as dist
import torch
from deepspeed.accelerator import get_accelerator
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

from transformers import HfArgumentParser, Seq2SeqTrainingArguments
import deepspeed
from transformers.deepspeed import HfDeepSpeedConfig

from accelerate import Accelerator

class Loader(Dataset):
    def __init__(self):
        ...
    def __len__(self):
        return 1
    def __getitem__(self, i):
        return "[INST] You are a helpy helper code writer, help the developer continue with this program[\INST] def main():\n    print(\"Hello W"
@dataclass
class Arguments:
    text: str = field(metadata="A field", default="Hello W")

def main():
    deepspeed.init_distributed()
    accelerator = get_accelerator()
    print(accelerator.__dict__)
    model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
    config = AutoConfig.from_pretrained(model_name)
    config.num_experts_per_tok= 8
    print(config)
    deepspeed.init_distributed("nccl")
    
    ds_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True,
        },
        "zero_optimization": {
            "stage": 3,
            "stage3_prefetch_bucket_size": 4 * config.hidden_size * config.hidden_size, # 0, 
            "stage3_param_persistence_threshold": config.hidden_size,
            "stage3_max_live_parameters": 4 * config.hidden_size * config.hidden_size,
            "offload_param": {"device": "cpu", "pin_memory": True}
        },
        "steps_per_print": 2000,
        "train_batch_size": 1,
        "train_micro_batch_size_per_gpu": 1,
        "wall_clock_breakdown": False,
    }
    dschf = HfDeepSpeedConfig(
        ds_config
    ) 

    parser = HfArgumentParser((Seq2SeqTrainingArguments,))
    loader = DataLoader(dataset=Loader(), batch_size=1)
    # args = parser.parse_args_into_dataclasses()[0]
    accelerator.empty_cache()
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
    
    
    
    model = deepspeed.init_inference(model, ds_config)[0]
    # model = ds.module
    msg = "[INST] You are a helpy helper code writer, help the developer continue with this program[\INST] def main():\n    print(\"Hello W"
   
    for prompt in loader:
        # if torch.distributed.get_rank() == 0:
        tokens = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")
        print(tokens)
        print(tokenizer.batch_decode(model.generate(tokens, max_new_tokens=100)))

if __name__=="__main__":

    main()

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.

The issue names ../demo.py, config-geforfuerte02.yaml, and the accelerate launch command; start by reproducing that command with the supplied Mixtral and DeepSpeed configuration. Done means establishing whether CPU-offloaded inference across the two computers is supported and, if so, documenting a working configuration and expected GPU usage or performance.