Continued Pre-Training and Mid-Training

Nova Forge CPT offers advanced capabilities beyond standard CPT, including access to intermediate checkpoints and data mixing with Nova's pre-training corpus. These features enable more efficient domain adaptation and better preservation of the model's general capabilities.

What are intermediate checkpoints and why are they needed?

Intermediate checkpoints are snapshots of the Amazon Nova model saved at different stages of the pre-training, before the model reaches its final production-ready state. During model development, Amazon Nova undergoes multiple training phases: initial pre-training with constant learning rate, learning rate ramp-down, context extension training, and finally instruction-following alignment and safety training. For CPT, intermediate checkpoints are often preferable to the final Prod checkpoint because they are more plastic and receptive to domain adaptation. The Prod checkpoint has undergone extensive instruction-following alignment and safety training, which optimizes the model for general conversational use but can make it resistant to learning new domain-specific patterns during CPT. In contrast, Partially and Fully pre-trained text only checkpoints retain the model's pre-training characteristics. They haven't been heavily steered toward specific behaviors, making them more efficient starting points for domain adaptation. When performing large-scale CPT (>10B tokens), starting from intermediate checkpoints typically results in faster convergence, better training stability, and more effective domain knowledge acquisition. However, for small-scale CPT (<10B tokens), or when instruction-following capabilities need to be preserved, the Prod checkpoint may be more appropriate as it allows domain adaptation while maintaining the model's conversational abilities.

Multiple intermediate checkpoints are necessary for CPT because they offer different levels of model plasticity that affect how efficiently the model can absorb new domain knowledge. The final Prod checkpoint has undergone extensive instruction-following alignment and safety training, which optimizes it for general conversational use but makes it resistant to learning new domain-specific patterns. In other words, It has been hardened through post-training. In contrast, earlier checkpoints retain the model's pre-training characteristics and haven't been heavily steered toward specific behaviors, making them more plastic and receptive to domain adaptation.

To achieve the best training efficiency, multiple intermediate checkpoints are provided.

What checkpoints are available?

Nova 1.0

Amazon Nova 1.0 family has three models (Micro, Lite, Pro) and for each model there are three checkpoints available.

Nova 2.0

There are three Amazon Nova Lite 2.0 checkpoints.

The following table elaborates on the different conditions for pre- and mid-training.

Which checkpoint to use?

Partially pre-trained text only and fully pre-trained text only checkpoints typically converge faster and require fewer training steps for domain adaptation. However, they have no instruction tuning and would need to undergo post training steps to be able to perform useful tasks and follow instructions. GA checkpoint may require more steps to adapt but provides safer starting point for small-scale experiments and will maintain some of it post training capabilities even after CPT training.

In general, with large training datasets (>10B tokens), start from partially pre-trained text only or fully pre-trained text only checkpoints for more efficient and stable training, as the model's knowledge base will be substantially modified. With small datasets (<10B tokens), use the GA checkpoint to preserve instruction-following capabilities while adapting to the domain.

How to use data mixing for 1.0 or 2.0 models?

When performing CPT with a new domain data, it is highly beneficial to mix the new data with some of the data used previously in the pre-training stage of the model. Mixing old data with new domain data solves two problems:

Nova CPT Data Mixing Capabilities

Access to Amazon Nova pre-training data and checkpoints is one of the core offerings of the Amazon Nova CPT customization. Amazon Nova CPT customization enables easy mixing of domain data with Amazon Nova's pre-training corpus. Further, the sampling ratio of the specific Amazon Nova data categories (e.g., code, math, reasoning, etc) can be changed and their proportions controlled to complement domain data. This allows reinforcement of capabilities that align with the use case while adapting the model to the specific domain.

Finding the Optimal Mixing Ratio

The optimal ratio of Amazon Nova data versus domain data depends on the dataset's domain, complexity, size, quality, and the importance of maintaining general capabilities. This ratio must be discovered through experimentation. An experiment framework to decide on how much Amazon Nova data to mix is as follows.

Select a representative subset of domain data (e.g., 5B tokens) and keep this constant across all experimental runs.

Run small-scale CPT experiments varying only the amount of Amazon Nova data mixed in:

Evaluate each checkpoint on in domain and general domain benchmarks. Also evaluate the starting checkpoint (Amazon Nova checkpoint before any training).

Once the optimal mixing ratio has been identified, run full-scale CPT using the complete domain dataset with the selected mixing ratio.

Dissecting the Data Mixing Categories

Below we dissect each available category in Data Mixing, for you to make best decision of what data categories makes most sense to be represented in your overall data mixture.

How to Enable Data Mixing

Add the data_mixing section to your recipe with the appropriate percentage distribution across dataset categories. The nova_data percentages must sum to 100.

Nova 1.0 Configuration with Data mixing

run:
  name: "cpt-job-name"             # A descriptive name for your training job
  model_type: "amazon.nova-lite-v1:0:300k"  # Model variant specification, do not change
  model_name_or_path: "nova-lite/prod"
  replicas: 4 
  data_s3_path: "s3://path/to/data/xyz.jsonl"
  output_s3_path: "s3://path/to/output/checkpoint"
  
skip_recipe_validation: true
training_config:
  max_length: 32768
  global_batch_size: 64
  trainer:
    max_steps: 5000
  model:
    hidden_dropout: 0.1
    attention_dropout: 0.1
    ffn_dropout: 0.1
    optim:
      lr: 1.5e-05
      name: distributed_fused_adam
      adam_w_mode: true
      eps: 1.0e-06
      weight_decay: 0.05
      betas:
      - 0.9
      - 0.999
      sched:
        warmup_steps: 500
        constant_steps: 0
        min_lr: 1.5e-06
        
data_mixing:
  dataset_catalog: cpt_text_lite
  sources:
    nova_data:
      en-entertainment: 0.11%
      en-factual: 4.83%
      en-legal: 0.48%
      en-long-form-text: 6.26%
      en-mined: 16.79%
      en-other: 1.79%
      en-scientific: 10.53%
      en-social: 12.43%
      en-techqa: 13.95%
      code: 7.50%
      high-util-lang: 8.05%
      low-util-lang: 6.51%
      math: 8.76%
      en-finance: 1%
      tables: 1%
    customer_data:
      percent: 90

What do these categories mean

Nova 2.0 Configuration with data mixing

# Note:
# This recipe can run on p5.48xlarge

# Run config
display_name: "Nova Lite Pretrain on P5 GPU"
versions: ["2.0"]
instance_types: ["ml.p5.48xlarge"]

run:
  name: "my-cpt-run"     # A descriptive name for your training job
  model_type: "amazon.nova-2-lite-v1:0:256k" # Model variant specification, do not change
  model_name_or_path: "nova-lite-2/prod" # Base model path, do not change
  replicas: 8       # Number of compute instances for training, allowed values are 4, 8, 16, 32
  data_s3_path: ""       # Customer data paths
  validation_data_s3_path: ""        # Customer validation data paths
  output_s3_path: ""   # Output artifact path, SageMaker HyperPod job-specific configuration - not compatible with standard SageMaker Training jobs

## Training specific configs
training_config:
  task_type: cpt
  max_length: 8192              # Maximum context window size (tokens)
  global_batch_size: 64        # Global batch size, allowed values are 32, 64, 128, 256.

  trainer:
    max_steps: 10               # The number of training steps to run total
    val_check_interval: 10      # The number of steps between running validation
    limit_val_batches: 2        # Batches of the validation set to use each trigger

  model:
    hidden_dropout: 0.0           # Dropout for hidden states, must be between 0.0 and 1.0
    attention_dropout: 0.0        # Dropout for attention weights, must be between 0.0 and 1.0

  optim:
    optimizer: adam
    lr: 1e-5                      # Learning rate
    name: distributed_fused_adam  # Optimizer algorithm, do not change
    adam_w_mode: true             # Enable AdamW mode
    eps: 1e-06                    # Epsilon for numerical stability
    weight_decay: 0.0             # L2 regularization strength, must be between 0.0 and 1.0
    adam_beta1: 0.9               # Beta1 for Adam optimizer
    adam_beta2: 0.95              # Beta2 for Adam optimizer
    sched:
      warmup_steps: 10            # Learning rate warmup steps
      constant_steps: 0           # Steps at constant learning rate
      min_lr: 1e-6                # Minimum learning rate, must be lower than lr

data_mixing:
  dataset_catalog: cpt_text_lite
  sources:
    nova_data:   # percent inputs for Nova data must sum to 100%; use 0% if you want to exclude a data grouping
      agents: 20
      business-and-finance: 4
      scientific: 10
      code: 5
      factual-and-news: 5
      longform-text: 6
      health-and-medicine: 1
      humanities-and-education: 1
      legal: 1
      math: 9
      additional-languages: 15
      social-and-personal-interest: 11
      entertainment: 0.5
      reasoning: 10
      other: 0.5
      tables: 1
    customer_data: # percent input of customer data. 100 = use only customer data, 0 = use only the nova_data mix above
      percent: 25

What do these categories mean

Note: Nova 2.0 includes additional reasoning-specific categories (e.g., reasoning-code, reasoning-math, reasoning-instruction-following) that are not available in Nova 1.0.

Summary of Categories and Info Labels:

Parameter Guide

The total number of tokens used in training can be calculated from max_length * global_batch_size * max_steps

Limitations

Current CPT only supports text data and does not support any customer multi-modal datasets.