Resolve DDP Deadlock on Multi-GPU Runtimes by Disabling device_map='auto'

[purvanshjoshi]

Technical Overview

Using device_map="auto" when launching training on multi-GPU nodes with accelerate launch causes Hugging Face's pipeline to automatically slice model parameters across all visible GPUs. This contradicts Distributed Data Parallel (DDP) fundamentals, where each GPU must host a complete, isolated copy of the model parameters. This conflict results in severe communication stalls, process hangs, and device deadlocks.

Affected Modules

• File: src/train.py (Line 32-56)
• Target Classes: WhisperForConditionalGeneration.from_pretrained()

Detailed Traceback / Context

Distributed training processes freeze indefinitely at trainer.train() on step 1 without logging errors or utilizing system resources.

Acceptance Criteria

1. Isolated GPU Contexts: Each subprocess allocates itself strictly to one dedicated GPU handled by the system's local rank.
2. Successful Multi-GPU Scaling: Effective training step latency drops proportionally when scaling from a single T4 to dual-T4 GPUs.
3. No Process Deadlocks: Zero synchronization stalls or hanging processes on DDP initialization.

Proposed Implementation Approach

Enforce localized rank allocations by disabling automatic mapping and hardcoding the local rank ID from process environment variables:

local_rank = int(os.environ.get("LOCAL_RANK", 0))
device = f"cuda:{local_rank}"

model = WhisperForConditionalGeneration.from_pretrained(
    args.model_name, 
    load_in_8bit=True,
    device_map={"": device} # Pin explicitly to the active process GPU rank
)

Severity & Priority

• Severity: High (Blocks distributed dual-GPU scaling)
• Priority: P1

[Saksham-dev-del]

Hi @purvanshjoshi,

I would like to work on this issue. I will update the model loading call in src/train.py (lines 32–56) to replace device_map="auto" with an explicit local-rank pinning approach, reading LOCAL_RANK from the environment and mapping the model to cuda:{local_rank} per process.

I will validate this on a dual-GPU setup to confirm training no longer stalls at trainer.train() and that step latency scales appropriately across GPUs.

Could you please assign this issue to me? I will submit a PR for review once complete.