# Distributed training with uneven inputs

**URL:** <https://discuss.ray.io/t/distributed-training-with-uneven-inputs/12427>\
**Category:** Ray Train\
**Created:** [October 12, 2023, 8:35am UTC](https://discuss.ray.io/t/distributed-training-with-uneven-inputs/12427 "2023-10-12T08:35:11Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![ben1](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/ben1/32/5259_2.png) [@ben1](https://discuss.ray.io/u/ben1)\
**Post date:** [October 12, 2023, 8:35am UTC](https://discuss.ray.io/t/distributed-training-with-uneven-inputs/12427/1 "2023-10-12T08:35:11Z")

</div>

I’m trying to replace torch.DDP with ray.train, and hit the problem of different DDP instance having different number of input batches.

DDP works with the ‘Join’ context manager can fix this problem, as shown below.  
How to handle this problem using ray.train?

```auto
...
    model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank])
    with model.join():
        for input in inputs:
            loss = model(input).sum()
            ...

```

[https://pytorch.org/tutorials/advanced/generic\_join.html#using-join-with-distributeddataparallel](https://pytorch.org/tutorials/advanced/generic_join.html#using-join-with-distributeddataparallel)

---

<div class="post-metadata">

**Author:** ![matthewdeng](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/matthewdeng/32/1446_2.png) [@matthewdeng](https://discuss.ray.io/u/matthewdeng)\
**Post date:** [October 13, 2023, 11:36pm UTC](https://discuss.ray.io/t/distributed-training-with-uneven-inputs/12427/2 "2023-10-13T23:36:18Z")

</div>

Hey @ben1, Ray Train supports Torch DDP and you should be able to run this as well.

---

<div class="post-metadata">

**Author:** ![ben1](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/ben1/32/5259_2.png) [@ben1](https://discuss.ray.io/u/ben1)\
**Post date:** [October 17, 2023, 12:50pm UTC](https://discuss.ray.io/t/distributed-training-with-uneven-inputs/12427/3 "2023-10-17T12:50:32Z")

</div>

Yes, I find out that the actually blocking part is train.report.  
Due to the uneven inputs, some worker finish the dataloader loop early. But train.report must be called the same number of times by all workers.

How to handle this problem using ray.train?

```auto
with model.join():
        iter_idx = 0
        for input in inputs:
            loss = model(input).sum()
            iter_idx += 1
            if iter_idx % report_num == 0:
                train.report(metrics)
            ...

```

---

<div class="post-metadata">

**Author:** ![justinvyu](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/justinvyu/32/4309_2.png) [@justinvyu](https://discuss.ray.io/u/justinvyu)\
**Post date:** [October 26, 2023, 5:17pm UTC](https://discuss.ray.io/t/distributed-training-with-uneven-inputs/12427/4 "2023-10-26T17:17:19Z")

</div>

@ben1 It may be easier to just report one set of aggregated metrics at the end of the epoch, rather than every N steps within the loop.

```python
for epoch in range(epochs):
    for input in inputs:  
        ...
    if epoch % report_num == 0:
        train.report(aggregated_metrics)

```

`train.report` does act as a barrier to maintain a consistent metric/checkpoint report iteration across workers.