Added training guidlines on how to run DLRM/HSTU on v6

Author: ajkv-google

training.md

@@ -0,0 +1,86 @@
+# Model Training Guide
+
+This guide explains how to set up the environment and train the HSTU/DLRM models on Cloud TPU v6 (or other TPU generations).
+
+## Option 1: Virtual Environment (Recommended for Dev)
+
+If you are developing on a TPU VM directly, use a virtual environment to avoid conflicts with the system-level Python packages.
+
+#### 1. Prerequisites
+Ensure you have **Python 3.11+** installed.
+```bash
+python3 --version
+```
+#### Create the venv\
+```bash
+python3 -m venv venv
+```
+#### Activate it
+```bash
+source venv/bin/activate
+```
+
+### 2. Create and Activate Virtual Environment
+Run the following from the root of the repository:
+```bash
+# Create the venv
+python3 -m venv venv
+
+# Activate it
+source venv/bin/activate
+```
+
+### 3. Install Dependencies
+```bash
+pip install -r requirements.txt
+```
+We need to force a specific version of Protobuf to ensure compatibility with our TPU stack. Run this exactly as shown:
+```bash
+pip install "protobuf>=6.31.1" --no-deps
+```
+The --no-deps flag is required to prevent pip from downgrading it due to strict dependency pinning in other libraries.
+
+### 4. Run the Training for DLRM
+```bash
+python dlrm_experiment_test.py
+```
+
+## Option 2: Docker (Recommended for Production)
+
+If you prefer not to manage a virtual environment or want to deploy this as a container, you can build a Docker image.
+
+## 1. Build the Image
+reate a file named `Dockerfile` in the root of the repository:
+
+```dockerfile
+# Use an official Python 3.11 runtime as a parent image
+FROM python:3.11-slim
+
+# Set the working directory
+WORKDIR /app
+
+# Copy the current directory contents into the container
+COPY . /app
+
+# Install system tools if needed (e.g., git)
+RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
+
+# Install dependencies
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
+
+# Force install the specific protobuf version
+RUN pip install "protobuf>=6.31.1" --no-deps
+
+# Default command to run the training script
+CMD ["python", "recml/examples/dlrm_experiment_test.py"]
+```
+
+You can use this dockerfile to run the DLRM model experiment from this repo in your own environment. 
+
+
+
+
+
+
+