Why buffers are not synchronized in zero3?

[kiradiso]

Hellow, i launch my model by accelerate launch with zero3 and found all parameters are partitioned. For my context, i apply processor.vision_tokenizer only for inference which is conducted by following code:

with deepspeed.zero.GatheredParameters(processor.vision_tokenizer.parameters(), modifier_rank=None):
    inputs = processor.vision_tokenizer(pixel_values)

However, i found inputs varies from different ranks (with same pixel_values). After some checks, i found the buffers of nn.BatchNorm3d are different (runner_vars/means). Only the buffers of rank 0 is correct.
So my questions are:

1. how to make these buffers consistent across ranks?
2. Whether exists an more better way to conduct inference for processor.vision_tokenizer, which is invariant to the transformer for training.

[tjruwase]

@kiradiso, thanks for your questions

1. Try setting modifier_rank=0 to get consistent buffers. See here.
2. Please clarify what you mean by better? Do you mean low latency, higher throughput, memory efficient, etc.

[kiradiso]

Thanks for your answer

1. I find that setting modifier_rank=0 is not worked for me. My code is like this:
   Program is launched by accelerate launch --config_file={a zero3 config} ...
   In main program, i have：

accelerator = Accelerator(
        gradient_accumulation_steps=1,
        mixed_precision="bf16",
        log_with="wandb",
        project_dir={a logging_dir},
        split_batches=True,
    )
image_tokenizer = AutoModel.from_pretrained("BAAI/Emu3-VisionTokenizer", trust_remote_code=True)
with deepspeed.zero.GatheredParameters(image_tokenizer.parameters(), modifier_rank=0):
    print(image_tokenizer.encoder.time_res_stack[0].norm1._buffers, torch.distributed.get_rank())

The outputs are like this:

OrderedDict([('running_mean', tensor([ 8.1875,  6.7812,  5.2812, 11.0625], device='cuda:0', dtype=torch.bfloat16)), 
('running_var', tensor([4.5312, 1.6484, 2.1406, 5.8438], device='cuda:0', dtype=torch.bfloat16)), 
('num_batches_tracked', tensor(650000, device='cuda:0'))]) 0

OrderedDict([('running_mean', tensor([0., 0., 0., 0.], device='cuda:1', dtype=torch.bfloat16)), 
('running_var', tensor([1., 1., 1., 1.], device='cuda:1', dtype=torch.bfloat16)), 
('num_batches_tracked', tensor(0, device='cuda:1'))]) 1

One solution is moving the definition of image_tokenizer before the definition of accelerator, then, image_tokenizer is not partitioned.

2. I don't want image_tokenizer to be partitioned, but with my usage, any model defined after the definition of accelerator is partitioned. If i don't use accelerator, i think i can use deepspeed.zero.Init as a context manager to actively restrict the scope of auto-partition (i have not tried).

[tjruwase]

I think the problem is caused by underlying HfDeepSpeedConfig used for managing the DeepSpeed integration. When used in a global context, this imposes the first ds_config on all DeepSpeed interactions. One workaround is to move each model or tokenizer initialization into separate functions to restrict the HfDeepSpeedConfig to function scope instead of global. I realize that Accelerate does not expose HfDeepSpeedConfig but I think the backend operates similarly.

If you are unable to share your full code with the Accelerate call context, then perhaps the following transformer example could explain more: https://gist.github.com/tjruwase/f1adccf087b8fa269ffce2ab91c4f1c6