Model with balanced load runs slower than the imbalanced

[CanyonWind]

Hi, sorry for spamming the issues. Another thing I noticed when using scattermoe is that, as title says, the model enforced with load balancing runs lightly slower than the model with imbalanced load. Could you please sharing any insights whether that should be expected?

The speed difference isn't significant but fairly consistent -- model A is always 2-3% slower. In large scale training, minor speed difference still matters a lot. Would appreciate any thought on this.

Two identical models:

• Model-A with load balancing loss and the load looks like below. As you can see the load are pretty evenly distributed

  • Speed: 1.17 iter/s

    Iter 38350, [[628, 649, 658, 630, 668, 652, 634, 601], [654, 627, 632, 676, 586, 646, 647, 652], [643, 633, 660, 604, 634, 664, 647, 635], [613, 640, 656, 606, 659, 641, 648, 657], [609, 624, 627, 658, 587, 705, 624, 686], [653, 615, 679, 646, 605, 615, 641, 666], [702, 680, 599, 675, 622, 620, 605, 617], [711, 658, 672, 623, 590, 688, 609, 569], [644, 629, 632, 627, 652, 639, 645, 652], [644, 659, 585, 701, 595, 648, 618, 670], [781, 548, 647, 600, 636, 566, 713, 629], [601, 617, 667, 674, 602, 675, 596, 688], [642, 692, 647, 603, 690, 621, 614, 611], [684, 613, 667, 696, 580, 632, 634, 614], [622, 632, 593, 620, 607, 718, 726, 602], [595, 682, 652, 702, 613, 640, 668, 568], [615, 577, 662, 686, 640, 661, 607, 672], [609, 619, 607, 580, 661, 722, 671, 651], [687, 676, 597, 660, 612, 617, 659, 612], [617, 585, 645, 667, 697, 630, 639, 640], [633, 619, 646, 620, 706, 639, 662, 595]]

• Model-B without load balancing. The experts are assigned imbalancedly

  • Speed: 1.20 iter/s

    Iter 38350, [[2559, 0, 0, 0, 0, 0, 1, 2560], [2550, 98, 206, 111, 37, 79, 1679, 360], [358, 846, 317, 1601, 213, 111, 1154, 520], [24, 2187, 2538, 59, 45, 180, 33, 54], [345, 414, 933, 937, 12, 1350, 1076, 53], [19, 505, 521, 2467, 117, 131, 1234, 126], [558, 1107, 83, 489, 2560, 11, 283, 29], [2410, 419, 125, 12, 536, 946, 184, 488], [579, 2197, 667, 0, 70, 424, 358, 825], [142, 59, 394, 28, 1859, 67, 302, 2269], [433, 1317, 211, 364, 593, 2063, 69, 70], [6, 1798, 683, 12, 3, 2041, 120, 457], [232, 258, 966, 355, 539, 79, 2559, 132], [455, 135, 2554, 567, 479, 68, 628, 234], [13, 169, 2186, 550, 0, 22, 2, 2178], [5, 22, 798, 4, 285, 2553, 1436, 17], [1492, 2137, 363, 1082, 17, 22, 1, 6], [0, 0, 94, 2558, 2462, 6, 0, 0], [0, 3, 1, 2, 0, 2555, 0, 2559], [636, 831, 0, 23, 0, 0, 1734, 1896], [850, 0, 0, 2234, 2036, 0, 0, 0]]

[shawntan]

Unsure, but I would say this seems expected, especially if you can skip the computation of some experts when no tokens are routed to it.

Generally better to have load balancing on, otherwise you could be severely underutilising some experts, or not using them at all.

[CanyonWind]

especially if you can skip the computation of some experts when no tokens are routed to it.

Thanks, but as scattermoe doesn't set expert capacity, shouldn't the computation cost theoretically always be the same no matter the load is balanced or not? As long as number of total tokens and expert dimensions are the same, scattermoe seems to have no difference regarding to these two scenarios' total computation.

I''m not questioning whether turning on lb, of which it should be for better quality; but wondering why the actual speed shows the difference.

[shawntan]

As far as my understanding of GPU programming goes, there can likely be differences due to caching. As a result, reusing an expert should be faster because the weight block is already in the cache.

Without more information that'd be my guess for where the 'speedup' might be coming from.