Parallelize neuron training processes for each neuron core

[mselim00]

Enables full multi-processing across all neuron cores, and corrects an earlier issue where world size wasn't being correctly determined (i.e., each process was in its own process group). Changes from mpirun to torchrun.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[mselim00]

Will publish test results in a bit

[mselim00]

Ran the test on 2 nodes locally with go test

2025/01/27 19:50:34 Parsed throughput from 56 ranks. Total=2798.46 samples/s, Average=49.97 samples/s
2025/01/27 19:50:34 Average Throughput: 49.97 samples/second
2025/01/27 19:50:34 Parsed average epoch time from 56 ranks. Sum=17.36s, Average=0.31s
--- PASS: TestBertTraining (675.16s)
    --- PASS: TestBertTraining/bert-training (675.16s)
        --- PASS: TestBertTraining/bert-training/Neuron_training_Job_succeeds (675.14s)
PASS

There's some issue collecting the throughput info. This run showed 56 ranks, others show some other random number. I thought this was a RegEx issue, which I've fixed, but we still see the problem. Might work on this separately though, the ranks parsed from are probably representative of the group, and afaict nvidia training currently only parses from the master proc.

[mattcjo]

Is explicit service creation required for torchrun?

[mselim00]

yeah this service is required so we can dynamically determine the master node's IP with bert-training-0.training in the job spec

[mattcjo]

Ran the test on 2 nodes locally with go test

2025/01/27 19:50:34 Parsed throughput from 56 ranks. Total=2798.46 samples/s, Average=49.97 samples/s
2025/01/27 19:50:34 Average Throughput: 49.97 samples/second
2025/01/27 19:50:34 Parsed average epoch time from 56 ranks. Sum=17.36s, Average=0.31s
--- PASS: TestBertTraining (675.16s)
    --- PASS: TestBertTraining/bert-training (675.16s)
        --- PASS: TestBertTraining/bert-training/Neuron_training_Job_succeeds (675.14s)
PASS

There's some issue collecting the throughput info. This run showed 56 ranks, others show some other random number. I thought this was a RegEx issue, which I've fixed, but we still see the problem. Might work on this separately though, the ranks parsed from are probably representative of the group, and afaict nvidia training currently only parses from the master proc.

@mselim00 This is slightly concerning. Are you able to confirm expected number of processes is running even if metrics seem off?

[mselim00]

Ran the test on 2 nodes locally with go test

2025/01/27 19:50:34 Parsed throughput from 56 ranks. Total=2798.46 samples/s, Average=49.97 samples/s
2025/01/27 19:50:34 Average Throughput: 49.97 samples/second
2025/01/27 19:50:34 Parsed average epoch time from 56 ranks. Sum=17.36s, Average=0.31s
--- PASS: TestBertTraining (675.16s)
    --- PASS: TestBertTraining/bert-training (675.16s)
        --- PASS: TestBertTraining/bert-training/Neuron_training_Job_succeeds (675.14s)
PASS

There's some issue collecting the throughput info. This run showed 56 ranks, others show some other random number. I thought this was a RegEx issue, which I've fixed, but we still see the problem. Might work on this separately though, the ranks parsed from are probably representative of the group, and afaict nvidia training currently only parses from the master proc.

@mselim00 This is slightly concerning. Are you able to confirm expected number of processes is running even if metrics seem off?

Yep, I manually checked that we have logs from all 64 ranks, that all of them print those metrics, and that all of them print the training complete log line. I'm not sure as to the root cause atm, just know that it's probably not a RegEx issue at this point.

[mattcjo]

LGTM. Approving since CI check failure is unrelated. Merge once fixed.