Use AWS Neuron sdk 2.21

[dacorvo]

What does this PR do?

This bumps the AWS Neuron SDK version to use AWS Neuron SDK 2.21, which is the first SDK compatible with trn2 instances.

The underlying pytorch version is now 2.5.1, which implies significant changes in the XLA stack.

This leads to compilation errors in:

• training tests,

Invoking RemoveOptimizationBarriers pass
2025-01-03 08:45:20.301338: F hilo/MLIRPasses/Transforms/VerifyAliasing.cc:125] false && "Read after write on aliased input/output buffer."

• stable diffusion tests.

***** Compiling unet *****
Using Neuron: --optlevel 2
2025-01-03T08:46:19Z [TEN404] (_convolution.158) Internal tensorizer error: RelaxPredicates:inst should be valid after relaxing predicates

• swin tests.

[TEN404] (_divide.1146) Internal tensorizer error: BirCodeGenLoop:Too many strides! {{{{0,+,1}[4],+,0}[2],+,4}[16],+,0}[2]

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[JingyaHuang]

@dacorvo For failing vision model tests, it seems that batch_size != 1 will all fail (in our tests, we applied batch_size = 2). I will open a ticket in neuron SDK repo, I could set batch size to 1 to make the CIs green, but I doubt if we should bump to that version...

[github-actions]

This PR is stale because it has been open 15 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[michaelbenayoun]

LGTM