Add tensor parallelism support for HF wrapper forward and lm_eval integration

[bigximik]

✨ Description

Add tensor parallelism support for HF wrapper forward and lm_eval integration

Closes #334

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

Key updates introduced in this PR:

1. Fixed a bug where the batch config was being read from the wrong place.
2. Added additional broadcast primitives and optimized _object_to_tensor for faster performance (following PyTorch sources).
3. Added tensor-parallel logits collection in the model head.
4. Added tensor-parallel support to forward.
5. Added coordinator-forward mode, which allows generate to run only on data-parallel leader ranks while tensor parallel workers participate through worker_forward.
6. Added model and pipeline parallelism support to the lm_eval wrapper.
7. Added wait barriers in critical places, as the standard 60s timeout on distributed primitives was insufficient in cases such as slow post-processing for some lm_eval tasks or when batches are incomplete and some data-parallel ranks have no data.

🗒️ Notes and Known Issues

• Manually tested with DP and TP on 2 GPUs, DP+TP on 4 GPUs, and also on a single GPU.
• There is a problem with CUDA memory fragmentation, potentially caused by scattering and broadcasting tensors of different size.
• For some tasks (e.g., Wikitext using sliding-window log-likelihood), processing is very slow with data and model parallel setup. This is likely due to logits being sent to rank 0 and offloaded to CPU before applying softmax. The problem is more severe with larger batch sizes.
• High memory usage was observed in general. For example, with the Qwen 1.5B model and batch size 3 per GPU, memory spikes to nearly 100% during evaluation.

[jlamypoirier]

timeout. Unnecessary long timeouts are often bad, so I recommend making it optional (default none) and enabling only as needed.

[bigximik]

Context

Conceptually, places like worker_forward or data-parallel_worker wait primitives should only exit under three conditions:

1. They receive work
2. They receive a finish message
3. The connection with peers/coordinator is lost (after some timeout)

However, this is not how torch.distributed works. It is designed for more or less synchronous communication, while here we are trying to adapt it for asynchronous communication.

Problem

If we set the default timeout to None, users will end up seeing random timeouts in different places.

Discussion

A better long-term solution would be to use a distributed messaging framework that is more appropriate for sending work and finish messages. However, introducing another communication layer into fast_llm is likely outside the scope of this PR.

Proposal

• Keep the default timeout as it is, applied only to these entry points. reset timeout after wait operation to default of 60 sec.
• Clarify the naming/description to avoid confusion.
• Add a TODO to revisit this later with a more suitable communication framework.

[jlamypoirier]

This doesn't seem relevant outside lm_eval. Any way to move it there?

[bigximik]

I initially thought about handling this differently, but since each subclass of the model has its own class, the only practical way I found was to use a dynamic class that constructs itself on the fly with type.
This lets us encapsulate forward of the fast_llm Hugging Face class and then pass it to generate.

Something like:

def wrap_hf_model(model):
    inner_forward = get_bounded_method(model.forward, model)
    wrapper_class = get_new_type(
        model.__class__,
        {
            "inner_forward": inner_forward,
            "forward": cordinator_forward,
            "worker_forward": worker_forward,
        },
    )
    model.__class__ = wrapper_class
    return model

Another option would be to create a static wrapper class, but that would require exposing and forwarding a lot of functionality that generate expects.

So instead, I decided to implement this in our HF wrapper, since it is implemented before any class specialization.

[jlamypoirier]

I'm not sure I'm following here. From what I understand, these methods are called by FastLLMLmEvalWrapper above which we are free to adjust as we want, and there isn't any dependency on the HF model so moving should be easy. Or are there some call in the. LM eval base class that absolutely enforce this structure?

This isn't present in typical HF models, so I'd prefer to avoid it.

[bigximik]

No, lm_eval calls FastLLMLmEvalWrapper, which then calls our model.generate which in turn makes multiple calls to forward, all runs entirely on TP rank 0. Our forward must be overridden to handle data distribution across all TP ranks.

----------------------------------------------------------------------------------------------------
|                                         Tensor Parallel Setup                                    |
----------------------------------------------------------------------------------------------------

   [ lm_eval ]                   [ forward_worker ]                [ forward_worker ]   ... more
        |                                 |                                  |
        v                                 v                                  v
+-------------------------+               |                                  |
| HF generate mixin       |               |                                  |
| model must  be HF       |               |                                  |
| (model.generate)        |               |                                  |
| - runs only on TP rank0 |               |                                  |
| - does multiple forward |               |                                  |
|   calls + sampling      |               |                                  |
|   beam search, etc.     |               |                                  |
+-------------------------+               |                                  |
        |                                 |                                  |
        v                                 |                                  |
+-------------------------+               |                                  |
| model.forward()         |---------------+----------------------------------+--> [wait for data]
| must be overridden      |               |                                  |     [long timeout]
| - orchestrates TP calls |               |                                  |     
+-------------------------+               |                                  |
        |                                 |                                  |
        v                                 v                                  v
+---------------------+       +---------------------+          +---------------------+   ... TP N-1
| TP rank 0           |       | TP rank 1           |          | TP rank 2           |
| model.forward_inner |       | model.forward_inner |          | model.forward_inner |
+---------------------+       +---------------------+          +---------------------+

Alternatives i have considered:

1. Composition (wrapper-model around HF specialization)

• Define a class OrchestratorModel that looks like an HF model (has forward, generate, etc.).
• It contains an inner HuggingfaceBaseModelForCausalLM (or subclass) that runs on workers.
• Orchestrator forward does the TP dispatch/gathering, then calls into the inner worker models as needed.
• generate (from HF) runs on this outer orchestrator class, which works because it just calls self.forward.

This is clean, explicit, stable — but involves boilerplate to replicate the HF interface.

2. Dynamic class injection (multiple inheritance / runtime patching)

• Build a class at runtime that inherits both:
  • HF specialization (HuggingfaceGPTModelForCausalLM, etc.)
  • OrchestratorMixin (overrides forward).
• Register that as the actual model class for the model object.
• generate is inherited unmodified from HF mixin, but calls OrchestratorMixin.forward.

This avoids extra wrapper code, but is “hackier” and could break with HF updates.

That’s why I dismissed the other options, but if we really want to keep HuggingfaceBaseModelForCausalLM completely unmodified, option 1 (composition) is likely the safer and more maintainable approach for our use case.