Cleaning up the Dockerfile after the ngc pytorch switch (#1170):

• Eliminate already installed apt packages
• sparse attn requirement lead to a triton downgrade
• flash attn is already part of the ngc container (in another version that is compatible with TE)

The flash-attn change might be required for TE, see the TE setup script. The following error from the docker log was indicating so, too. I am not 100% sure if the version upper limit is a technological requirement, or just a precaution. In any case this is now resolved, plus there's no need to build flash-attn.

#14 10.98 Installing collected packages: flash-attn
#14 10.98   Attempting uninstall: flash-attn
#14 10.98     Found existing installation: flash-attn 2.4.2
#14 10.98     Uninstalling flash-attn-2.4.2:
#14 11.74       Successfully uninstalled flash-attn-2.4.2
#14 13.81 ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
#14 13.81 transformer-engine 1.3.0 5b90b7f requires flash-attn!=2.0.9,!=2.1.0,<=2.4.2,>=2.0.6, but you have flash-attn 2.5.6 which is incompatible.
#14 13.81 Successfully installed flash-attn-2.5.6

Training pythia14M works (with gpt-j-residual=false , see #1174)
Training a 20B model with TP8 works as well

Just a few steps on 8xH100, not to convergence :)