Update the NCCL flags for A3+.

[yangyuwei]

End-to-end test passed: first created an A3+ cluster including 2 nodes and then run MaxText at 2 nodes.

[michelle-yooh]

Have you tested this on A3 nodes as well?

[michelle-yooh]

Could you help me understand how these flags were determined? Also what differences in results this change causes

[yangyuwei]

Have you tested this on A3 nodes as well?

No. It's not needed, given the changes in this PR only impact A3+, not A3. As you see it's under fastrak branch, FasTrak is specific for A3+, the A3 config is under tcpx branch and it remains the same.

Could you help me understand how these flags were determined? Also what differences in results this change causes

It's determined by setting USE_GPUDIRECT in the env vars, which is done via XPK (see https://screenshot.googleplex.com/6augXoEPxzCrPYZ) or via helm (https://screenshot.googleplex.com/59XtkdZKsN5g5Hm).

The changes are needed, in order to use the latest FasTrak release (Pony + Aperture), which has networking perf improvement. They are consistent with the internal config in google3.