Bugs in Distributed Loading of Non-Distributed Checkpoints

[bigximik]

🐞 Bugs in Distributed Loading of Non-Distributed Checkpoints and/or Model Creation via HF Wrapper

There are issues when loading non-distributed checkpoints or creating models via from_pretrained using the HF wrapper. (For details, see discussion below.)

• When tensor_parallel is not 1: model creation succeeds, but checkpoint loading fails.
• When sequence_parallel is not 1: model creation succeeds, but checkpoint loading fails.
• When pipeline_parallel is not 1: model creation fails at different points, depending on whether sequence_tensor_parallel is set.

[jlamypoirier]

🎯 Goal (What & Why)

Enable distributed loading of non-distributed checkpoint formats (in addition to the currently supported distributed checkpoint format). This is required to support distributed inference for models stored in non distributed checkpoint format like those in HF format. Currently we can only implement data parallel inference form those checkpoints.

I don't understand this feature request. It should already be possible to load non-distributed checkpoints in any format from a distributed setting. Also how does this relate to #199?

This is currently a blocker for #217, and by extension, for #199 — unless we decide to support only data parallelism for those features in the short term.

Notes:

1. As I understand, this is already a known limitation — see the TODOs at:

   • https://github.com/ServiceNow/Fast-LLM/blob/main/fast_llm/engine/multi_stage/fast_llm_model.py#L36

We could use more test cases, but afaik things are working.

• https://github.com/ServiceNow/Fast-LLM/blob/main/fast_llm/engine/multi_stage/stage_base.py#L319
  regarding distributed loading and shard slice imports.

I'm not sure where this comes from, I think it was a temporary todo that I forgot to remove. Imports definitely work.

2. How have we handled continued training for 70B-scale models so far? Have we only pre-trained large models from scratch?

State imports should work at basically any scale, so yes we support extended pretraining there.

[bigximik]

Thanks for the comment, @jlamypoirier.

I encountered errors when trying to load a checkpoint (e.g., Qwen) with distributed.tensor_parallel = 2.

The only difference from main branch is that I modified the from_pretrained function to accept config updates directly:

python -m torch.distributed.run --nproc-per-node=2 test_distributed.py

https://github.com/ServiceNow/Fast-LLM/blob/denis/generate/test_distributed.py

updates = {
      ("base_model", "transformer", "use_flash_attention"): attn_implementation is not None
      and attn_implementation == "flash_attention_2",
      ("distributed", "tensor_parallel"): 2,
      ("distributed", "pipeline_parallel"): 1,
      ("distributed", "sequence_data_parallel"): 1,
  }

  model_fm = HuggingfaceGPTModelForCausalLM.from_pretrained(
      CheckpointLoadConfig(
          path=checkpoint,
          format=Qwen2GPTHuggingfaceCheckpointFormat,
      ),
      updates,
  )

errors:

Global counter mismatch for parameter "layers.1.norm_1.weight" and shard "weights": 3072 != 1536
Global counter mismatch for parameter "layers.1.norm_2.weight" and shard "weights": 3072 != 1536
....

The model params seems to be dividable, so I assumed that distributed loading is not yet supported.

What do you think the problem might be here?

Thanks!

[bigximik]

#199 depends on generate, and to implement generate in tensor, pipeline, and sequence parallel modes, I need tests. However, I currently can't load a trained checkpoint in distributed mode (model distributed on several gpus) to run tests.

[jlamypoirier]

Global counter mismatch for parameter "layers.1.norm_1.weight" and shard "weights": 3072 != 1536
Global counter mismatch for parameter "layers.1.norm_2.weight" and shard "weights": 3072 != 1536
....

The model params seems to be dividable, so I assumed that distributed loading is not yet supported.

This should work, so there must be a bug somewhere. I'll give it a try later today, in the meantime could you try with ("distributed", "sequence_tensor_parallel"): True,? We want it anyway, and it's better supported.

[bigximik]

With sequence_tensor_parallel true still fails. Config:

("distributed", "sequence_tensor_parallel"): True
("distributed", "tensor_parallel"): 2 (others set to 1)

• Model is instantiated.
• Checkpoint loading fails.

---

Also i have tried several more configuration options:

Config: ("distributed", "pipeline_parallel"): 2 (others set to 1)

• If ("distributed", "sequence_tensor_parallel"): True is also set:
  Trying to set an implicit default for field sequence_tensor_parallel, but the field has already been set explicitly.

• If ("distributed", "sequence_tensor_parallel") is not set:

File "fast_llm/engine/multi_stage/fast_llm_model.py", line 93, in initialize_weights 
self._stages[tied_parameter.main_stage].weight_shard, 0, tied_parameter.group, timeout=timeout 
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AttributeError: 
'Stage' object has no attribute 'weight_shard'

---

Config: ("distributed", "sequence_data_parallel"): 2 (others set to 1)

• Model is created.
• Checkpoint loading fails with:
  Global counter mismatch for parameter "layers.0.word_embeddings_weight" and shard "weights": 466747392 != 233373696

[jlamypoirier]

These look like bugs. I'll have a look