Make Docker builds architecture-agnostic (x86_64 / aarch64) and modernize Dockerfile syntax

[azrael417]

Summary

Enables the TorchFort container images to build on both x86_64 and aarch64 (e.g. GB10 / DGX Spark, Grace) from the same Dockerfiles, driven by a single ARCH build-arg. Also modernizes ENV
syntax to silence BuildKit lint warnings, and relaxes one RL test tolerance that fails only on ARM due to floating-point accumulation differences.

Changes

Architecture-agnostic builds

• Added ARG ARCH=x86_64 to all three Dockerfiles. The value maps directly onto the arch strings used by the toolchains, so no translation table is needed:
  • NVHPC SDK paths: Linux_${ARCH} (Dockerfile — CUDA_HOME, CUPTI LD_LIBRARY_PATH, and the NCCL removal path).
  • HPCX tarball: hpcx-...-cuda13-${ARCH}.tbz (Dockerfile_gnu).
• build_docker.sh passes --build-arg ARCH=$(uname -m), which yields exactly x86_64 / aarch64 — matching both naming schemes. Native-host builds; the ARCH default keeps a plain docker build
  working unchanged.
• The CPU-only image (Dockerfile_gnu_cpuonly) has no arch-specific paths — the CPU PyTorch wheel resolves per-arch automatically — so it declares ARG ARCH only to consume the shared build-arg
  without a warning.

Dockerfile syntax modernization

• Converted all legacy ENV key value lines to ENV key=value (removes 7 LegacyKeyValueFormat warnings).
• Changed the first LD_LIBRARY_PATH assignment in each file to the ${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} form, so the separator/append only appears when the variable is already set. This
  removes the UndefinedVar warning on the bare-Ubuntu CPU image and keeps the GPU images robust regardless of whether the base image pre-sets the var (no trailing colon when unset).

Test fix (ARM)

• tests/rl/test_replay_buffer.cpp NStepConsistency: replaced EXPECT_FLOAT_EQ(rdiff, 0.) with EXPECT_NEAR(rdiff, 0., 1e-4). rdiff accumulates float32 rounding from the n-step discounted reward
  sum across the batch. The buffer and the test's manual recomputation use mathematically equivalent but differently-rounded paths (std::pow + torch tensor ops vs. repeated scalar float
  multiply), which diverge by a few ULP per term and differ across architectures (FMA contraction on aarch64). The buffer math is unchanged and correct on both arches; only the test's exact-zero
  comparison was too strict. sdiff (no arithmetic) stays on exact equality.

Testing

• Built and ran on aarch64 (single-GPU, DGX Spark): the CPU image builds warning-free; RL NStepConsistency now passes. The 9 skipped supervised tests are the multi-GPU (device-index-1) cases,
  expected on a single-GPU host — not a regression.
• x86_64 behavior is unchanged (default ARCH=x86_64, same numeric tolerances honored).

Notes for reviewers

• Cross-arch (emulated) builds are out of scope: uname -m reflects the build host, so this targets native builds. Docker's TARGETARCH would be needed for --platform emulation.
• Assumes the aarch64 HPCX tarball and the aarch64 NVHPC/CUDA base image variants exist at the referenced URLs/tags.

[romerojosh]

I pushed a modification to your Dockerfiles here to use TARGETARCH which is what is automatically exported as buildkit. This has two benefits:

1. No need to manually add the ARCH build argument, the architecture is automatically detected and accounted for.
2. This enables multi-arch builds with buildx.

There is an annoying wrinkle in that TARGETARCH values are amd64 and arm64 which don't correspond exactly the the expected ARCH values, but the build stage stuff at the top of the Dockerfiles handles this conditional setting. Apparently this is the only way to set ARCH conditionally, which is lame. Give the Dockerfiles a spin on your end to make sure they still work.