ltx_video.yml

#  Copyright 2025 Google LLC

#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at

#       https://www.apache.org/licenses/LICENSE-2.0

#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.


#hardware
hardware: 'tpu'
skip_jax_distributed_system: False

jax_cache_dir: ''
weights_dtype: 'bfloat16'
activations_dtype: 'bfloat16'


run_name: ''
output_dir: 'ltx-video-output'
save_config_to_gcs: False

#parallelism
mesh_axes: ['data', 'fsdp', 'tensor']
logical_axis_rules: [
                      ['batch', 'data'],
                      ['activation_heads', 'fsdp'],
                      ['activation_batch', ['data','fsdp']],
                      ['activation_kv', 'tensor'],
                      ['mlp','tensor'],
                      ['embed','fsdp'],
                      ['heads', 'tensor'],
                      ['norm', 'fsdp'],
                      ['conv_batch', ['data','fsdp']],
                      ['out_channels', 'tensor'],
                      ['conv_out', 'fsdp'],
                      ['conv_in', 'fsdp']
                    ]
data_sharding: [['data', 'fsdp', 'tensor']]
dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
dcn_fsdp_parallelism: -1
dcn_tensor_parallelism: 1
ici_data_parallelism: 1
ici_fsdp_parallelism: -1  # recommended ICI axis to be auto-sharded
ici_tensor_parallelism: 1




learning_rate_schedule_steps: -1
max_train_steps: 500 #TODO: change this
pretrained_model_name_or_path: ''
unet_checkpoint: ''
dataset_name: 'diffusers/pokemon-gpt4-captions'
train_split: 'train'
dataset_type: 'tf'
cache_latents_text_encoder_outputs: True
per_device_batch_size: 1
compile_topology_num_slices: -1 
quantization_local_shard_count: -1
jit_initializers: True 
enable_single_replica_ckpt_restoring: False