apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: jeesoo-rms-optim-v10
  namespace: kbm-g-np-motif
spec:
  managedBy: trainer.kubeflow.org/trainjob-controller
  podTemplateOverrides:
    - spec:
        containers:
          - name: node
            volumeMounts:
              - mountPath: /dev/shm
                name: shm
              - mountPath: /mair
                name: mair
        volumes:
          - emptyDir:
              medium: Memory
              sizeLimit: 64Gi
            name: shm
          - name: mair
            persistentVolumeClaim:
              claimName: mair
      targetJobs:
        - name: node
  runtimeRef:
    apiGroup: trainer.kubeflow.org
    kind: ClusterTrainingRuntime
    name: torch-distributed
  suspend: false
  trainer:
    args:
      - /bin/bash
      - '-c'
      - >
        ACTIVATIONPATH=/mair/team-sys/jeesoo/activation

        pip install triton matplotlib pandas

        echo "=== Building ==="

        cd $ACTIVATIONPATH

        pip uninstall -y activation 2>/dev/null; true
        rm -rf $ACTIVATIONPATH/build/temp.linux-x86_64-cpython-312 $ACTIVATIONPATH/_activation*.so $ACTIVATIONPATH/*.egg-info

        pip install --no-build-isolation --no-cache-dir -e . -v 2>&1 | tail -100

        python -c "import _activation; print('Build OK:', _activation)" || { echo "BUILD FAILED"; exit 0; }

        echo "=== Running RMS tests ==="

        cd $ACTIVATIONPATH

        python -m pytest tests/test_rms_norm.py -v 2>&1 | tail -40

        echo "=== Warmup ==="

        python -c "import torch; x=torch.randn(8192,1280,device='cuda',dtype=torch.bfloat16); [torch.mm(x.T,x) for _ in range(100)]; torch.cuda.synchronize(); print('warmup done')"

        echo "=== RMS Benchmark ==="

        cd $ACTIVATIONPATH/benchmarks

        python run_rms_bench.py 2>&1 | tee results/rms_optim_log.txt

        echo "=== Done ==="

        exit 0;
    env:
      - name: PYTHONUNBUFFERED
        value: '1'
      - name: PYTORCH_ALLOC_CONF
        value: expandable_segments:True
      - name: CUDA_LAUNCH_BLOCKING
        value: '0'
      - name: OMP_NUM_THREADS
        value: '1'
      - name: HF_HOME
        value: /mair/llm-dataset/hf_cache
    image: ghcr.io/motiftechnologies/llm-training:v0.1.3
    numNodes: 1
    numProcPerNode: 1
    resourcesPerNode:
      limits:
        cpu: '16'
        memory: 128Gi
        nvidia.com/gpu: '1'
      requests:
        cpu: '16'
        memory: 128Gi
        nvidia.com/gpu: '1'