setup training for 14b MoE

data/mixkit_i2v_full.json

@@ -0,0 +1,452 @@
+{
+  "data": [
+    {
+      "prompt": "Young man skating with a skateboard on the ramps with graffiti of a park with trees, on a sunny day.", 
+      "image_path": "images/mixkit-boy-skating-with-a-skateboard-in-a-park-with-ramps-34389.png"
+    },
+    {
+      "prompt": "In the midst of the joyous New Year's Eve celebration, the cheerful group of friends, their spirits lifted by the festivities, decides to immortalize the moment with a vibrant snapshot", 
+      "image_path": "images/mixkit-a-cheerful-group-of-friends-celebrate-new-years-eve-and-51525.png"
+    },
+    {
+      "prompt": "A man and a woman playing in a field with grass, during a bright afternoon, while cars pass by in the distance.", 
+      "image_path": "images/mixkit-a-cute-couple-playing-on-the-grass-4688.png"
+    },
+    {
+      "prompt": "Aerial view of a rocky mountain in the forest at a sunny day drone flight footage", 
+      "image_path": "images/mixkit-aerial-view-of-a-rocky-mountain-in-the-forest-50589.png"
+    },
+    {
+      "prompt": "A little girl wearing a pink security helmet and denim overall discovers the art of cycling amidst the serene park, as the camera captures her graceful progress.", 
+      "image_path": "images/mixkit-a-little-girl-cruises-through-the-forest-path-on-her-50088.png"
+    },
+    {
+      "prompt": "Aerial shot of a beach shore with sea waves. Big rocks on the sand at an alone beach.", 
+      "image_path": "images/mixkit-aerial-shot-of-a-beach-with-sea-waves-1087.png"
+    },
+    {
+      "prompt": "Young woman cleaning her house decorated with plants and decorations, while dancing happily to music in her headphones.", 
+      "image_path": "images/mixkit-woman-cleaning-her-house-dancing-happy-43379.png"
+    },
+    {
+      "prompt": "Aerial tour in a meadow surrounded by hills on the horizon, while some birds fly low over a lake.", 
+      "image_path": "images/mixkit-birds-flying-low-over-a-lake-in-a-meadow-41417.png"
+    },
+    {
+      "prompt": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.", 
+      "image_path": "images/mixkit-a-rancher-riding-a-horse-at-sunset-1143.png"
+    },
+    {
+      "prompt": "In the video, a martial artist dressed in a traditional white uniform with a black belt demonstrates a series of precise movements against a stark black background. The individual gracefully transitions between stances, embodying a sense of focused discipline and control. Each motion is executed with a deliberate pace, showcasing the fluidity of martial arts techniques. The soft lighting creates subtle highlights on the uniform, adding depth to the figure as it moves. The practitioner begins with an open-hand pose, feet firmly grounded, gradually shifting to a powerful forward punch. The fluidity of the sequence displays a mastery of balance and poise. Every trajectory of the limbs is precise and deliberate, capturing the elegance and strength of martial arts. The serene, isolated setting enhances the intensity and concentration of the practitioner. This visual presentation is an elegant interplay of motion and stillness, displaying the art form's discipline and grace.", 
+      "image_path": "images/mixkit-a-young-man-practicing-his-karate-moves-49635.png"
+    },
+    {
+      "prompt": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.", 
+      "image_path": "images/mixkit-small-group-of-people-doing-yoga-together-43730.png"
+    },
+    {
+      "prompt": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.", 
+      "image_path": "images/mixkit-dolphins-underwater-4133.png"
+    },
+    {
+      "prompt": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.", 
+      "image_path": "images/mixkit-skiers-on-a-snowy-slope-3327.png"
+    },
+    {
+      "prompt": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
+      "image_path": "images/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306.png"
+    },
+    {
+      "prompt": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
+      "image_path": "images/mixkit-curve-on-a-snowy-forest-road-3317.png"
+    },
+    {
+      "prompt": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
+      "image_path": "images/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996.png"
+    },
+    {
+      "prompt": "Far shot from high behind a beach on a sunny day, with crowds of people everywhere, buildings and hotels behind, and the sea on the horizon as ships cross it.",
+      "image_path": "480P_i2v_mixkit_images/Far shot from high behind a beach on a sunny day, with crowds of people everywhere, buildings and ho.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Skillful cyclist doing a wheelie on a bike while riding through a forest, on a dirt road, surrounded by many trees, in the morning.",
+      "image_path": "480P_i2v_mixkit_images/Skillful cyclist doing a wheelie on a bike while riding through a forest, on a dirt road, surrounded.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Portrait of a man with a mask and fluorescent paint dancing in front of the camera, under a blue party light.",
+      "image_path": "480P_i2v_mixkit_images/Portrait of a man with a mask and fluorescent paint dancing in front of the camera, under a blue par.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Timelapse of a luxury swimming pool in a boutique hotel.",
+      "image_path": "480P_i2v_mixkit_images/Timelapse of a luxury swimming pool in a boutique hotel..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Man dressed in 80's style dances very happily in his kitchen while listening to music on his radio and drinking wine.",
+      "image_path": "480P_i2v_mixkit_images/Man dressed in 80's style dances very happily in his kitchen while listening to music on his radio a.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Romantic couple embracing and looking at each other in the middle of a forest, during a break on a road trip through nature.",
+      "image_path": "480P_i2v_mixkit_images/Romantic couple embracing and looking at each other in the middle of a forest, during a break on a r.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Big city with many big buildings and a big capitol, in a view from a slow drive through the heights in the afternoon.",
+      "image_path": "480P_i2v_mixkit_images/Big city with many big buildings and a big capitol, in a view from a slow drive through the heights .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A man appears in different parts of an abandoned building, standing in the dark, swinging a baseball bat, outside.",
+      "image_path": "480P_i2v_mixkit_images/A man appears in different parts of an abandoned building, standing in the dark, swinging a baseball.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Many small lights out of focus yellow and blue and of different intensities, at night on a blue background.",
+      "image_path": "480P_i2v_mixkit_images/Many small lights out of focus yellow and blue and of different intensities, at night on a blue back.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Pair of jazz musicians performing a song with their saxophone and trombone on an abandoned train.",
+      "image_path": "480P_i2v_mixkit_images/Pair of jazz musicians performing a song with their saxophone and trombone on an abandoned train..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Pair of mountain bikers climbing a slope in the forest covered with rocks, with abundant trees and vegetation around.",
+      "image_path": "480P_i2v_mixkit_images/Pair of mountain bikers climbing a slope in the forest covered with rocks, with abundant trees and v.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A model with cyberpunk makeup poses in front of a neon light circle background.",
+      "image_path": "480P_i2v_mixkit_images/A model with cyberpunk makeup poses in front of a neon light circle background..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Couple of boy and girl walking hand in hand through a meadow among many cows that feed on grass and grass.",
+      "image_path": "480P_i2v_mixkit_images/Couple of boy and girl walking hand in hand through a meadow among many cows that feed on grass and .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Natural coastal relief on a sunny day, with large hills covered with dry grass in an arid climate, in a panoramic aerial view.",
+      "image_path": "480P_i2v_mixkit_images/Natural coastal relief on a sunny day, with large hills covered with dry grass in an arid climate, i.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Young woman lying on her bed, while chatting on her cell phone, in a close-up shot of her hands typing on the phone.",
+      "image_path": "480P_i2v_mixkit_images/Young woman lying on her bed, while chatting on her cell phone, in a close-up shot of her hands typi.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Meadow with a frozen lake in a forest of pines covered in snow and fog.",
+      "image_path": "480P_i2v_mixkit_images/Meadow with a frozen lake in a forest of pines covered in snow and fog..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Aerial shot of buses and cars on the freeway with a cloudy sky in the background.",
+      "image_path": "480P_i2v_mixkit_images/Aerial shot of buses and cars on the freeway with a cloudy sky in the background..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Beautiful countryside meadow during summer with trees, grass and little yellow flowers in the ground on a sunny day.",
+      "image_path": "480P_i2v_mixkit_images/Beautiful countryside meadow during summer with trees, grass and little yellow flowers in the ground.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A man wearing grey shorts jumps rope in a gym, weights and gym equipment in the background.",
+      "image_path": "480P_i2v_mixkit_images/A man wearing grey shorts jumps rope in a gym, weights and gym equipment in the background..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A young woman wearing a black mini skirt, a red sleeveless shirt and a red Santa hat stands on a stage and sings into a microphone. In the background, a guitar player.",
+      "image_path": "480P_i2v_mixkit_images/A young woman wearing a black mini skirt, a red sleeveless shirt and a red Santa hat stands on a sta.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Little baby with a pacifier, playing with a teddy bear, accompanied by his mom, both sitting on a bed, in a close view.",
+      "image_path": "480P_i2v_mixkit_images/Little baby with a pacifier, playing with a teddy bear, accompanied by his mom, both sitting on a be.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A saxophonist wearing a blazer dances while playing a song in a park.",
+      "image_path": "480P_i2v_mixkit_images/A saxophonist wearing a blazer dances while playing a song in a park..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A young woman with short hair wearing pink sunglasses chews gum and makes a bubble gum with the city in the background.",
+      "image_path": "480P_i2v_mixkit_images/A young woman with short hair wearing pink sunglasses chews gum and makes a bubble gum with the city.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "White flowers from a bridal bouquet, with the background out of focus.",
+      "image_path": "480P_i2v_mixkit_images/White flowers from a bridal bouquet, with the background out of focus..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Silhouette of a person looking at an image projected in front, with spectra and light patterns rotating around a center in the background.",
+      "image_path": "480P_i2v_mixkit_images/Silhouette of a person looking at an image projected in front, with spectra and light patterns rotat.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Natural aerial landscape with a relief covered with abundant trees and vegetation and a thick layer of mist.",
+      "image_path": "480P_i2v_mixkit_images/Natural aerial landscape with a relief covered with abundant trees and vegetation and a thick layer .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Boy dancing happily and cheerfully in a room set in the 80s.",
+      "image_path": "480P_i2v_mixkit_images/Boy dancing happily and cheerfully in a room set in the 80s..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Close up video of the legs of a young gymnast woman doing some acrobatics on a balance beam in a dark and empty gym.",
+      "image_path": "480P_i2v_mixkit_images/Close up video of the legs of a young gymnast woman doing some acrobatics on a balance beam in a dar.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Female scientist in a laboratory, with test tube, and more containers, looking through the lens of a microscope.",
+      "image_path": "480P_i2v_mixkit_images/Female scientist in a laboratory, with test tube, and more containers, looking through the lens of a.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Young woman getting makeup done by a professional makeup artist for a winter concept fashion photo shoot.",
+      "image_path": "480P_i2v_mixkit_images/Young woman getting makeup done by a professional makeup artist for a winter concept fashion photo s.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Loving couple sitting on a log on the shore of a lake outside, sharing an affectionate hug.",
+      "image_path": "480P_i2v_mixkit_images/Loving couple sitting on a log on the shore of a lake outside, sharing an affectionate hug..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Small family of father, mother and daughter, walking together through nature in a forest.",
+      "image_path": "480P_i2v_mixkit_images/Small family of father, mother and daughter, walking together through nature in a forest..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Flying over a peninsula covered in bushy trees, while discovering the sea around it, painted a beautiful turquoise blue, on a sunny day.",
+      "image_path": "480P_i2v_mixkit_images/Flying over a peninsula covered in bushy trees, while discovering the sea around it, painted a beaut.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A woman with long blonde hair wearing a white t-shirt and black bikini bottoms stands on a wooden deck smiling, with the beach and palm trees in the background.",
+      "image_path": "480P_i2v_mixkit_images/A woman with long blonde hair wearing a white t-shirt and black bikini bottoms stands on a wooden de.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "View towards the horizon in a flat desert at night, as the full moon glows brightly red.",
+      "image_path": "480P_i2v_mixkit_images/View towards the horizon in a flat desert at night, as the full moon glows brightly red..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "100 dollar bill thrown in the street seen up close with the background out of focus, while a person arrives, picks it up and walks away.",
+      "image_path": "480P_i2v_mixkit_images/100 dollar bill thrown in the street seen up close with the background out of focus, while a person .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Flying over a forest of abundant trees and vegetation, with a river and a dam, with houses and hills in the surroundings.",
+      "image_path": "480P_i2v_mixkit_images/Flying over a forest of abundant trees and vegetation, with a river and a dam, with houses and hills.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Artistic portrait of a man with his body and face stained with phosphorescent paint that glows with black light in a dark place.",
+      "image_path": "480P_i2v_mixkit_images/Artistic portrait of a man with his body and face stained with phosphorescent paint that glows with .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Aerial view of the calm blue sea made with a drone showing the horizon and the sea.",
+      "image_path": "480P_i2v_mixkit_images/Aerial view of the calm blue sea made with a drone showing the horizon and the sea..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Cars driving on a two way street underneath a tree, with traffic signs, buildings and the sun in the background.",
+      "image_path": "480P_i2v_mixkit_images/Cars driving on a two way street underneath a tree, with traffic signs, buildings and the sun in the.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Face of a young woman, close up portrait",
+      "image_path": "480P_i2v_mixkit_images/Face of a young woman, close up portrait.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Man with his body splattered with fluorescent paint while doing a slow conceptual dance, under a black light in a dark place with colored spotlights around.",
+      "image_path": "480P_i2v_mixkit_images/Man with his body splattered with fluorescent paint while doing a slow conceptual dance, under a bla.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Two women dressed in colorful retro styles, having fun taking photos in a room with lots of pastel colors.",
+      "image_path": "480P_i2v_mixkit_images/Two women dressed in colorful retro styles, having fun taking photos in a room with lots of pastel c.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "The silhouette of a person who moves his head reflects inside a city with tall buildings next to a lake.",
+      "image_path": "480P_i2v_mixkit_images/The silhouette of a person who moves his head reflects inside a city with tall buildings next to a l.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "View from a car window of the trees lining the road speeding by, with the dazzling sunlight in the sky, in a retro looking video.",
+      "image_path": "480P_i2v_mixkit_images/View from a car window of the trees lining the road speeding by, with the dazzling sunlight in the s.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Aerial view of two people on a tennis court playing tennis.",
+      "image_path": "480P_i2v_mixkit_images/Aerial view of two people on a tennis court playing tennis..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "View from the floor upwards of a young woman with long dark hair dancing slowly in green and purple neon lights",
+      "image_path": "480P_i2v_mixkit_images/View from the floor upwards of a young woman with long dark hair dancing slowly in green and purple .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Some friends dancing and having fun together in circles, at a party surrounded by colored lights at a party, in a fancy old place, in a view from below them.",
+      "image_path": "480P_i2v_mixkit_images/Some friends dancing and having fun together in circles, at a party surrounded by colored lights at .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    }
+  ]
+}

fastvideo/configs/models/dits/wanvideo.py

@@ -7,10 +7,12 @@ from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 def is_blocks(n: str, m) -> bool:
     return "blocks" in n and str.isdigit(n.split(".")[-1])
 
+def is_not_blocks(n: str, m) -> bool:
+    return "blocks" not in n and len(n.split(".")) == 1
 
 @dataclass
 class WanVideoArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks, is_not_blocks])
 
     param_names_mapping: dict = field(
         default_factory=lambda: {
@@ -99,7 +101,9 @@ class WanVideoArchConfig(DiTArchConfig):
     local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
     sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
     num_frames_per_block: int = 3
+    # sliding_window_num_frames: int = 19
     sliding_window_num_frames: int = 21
+    # sliding_window_num_frames: int = 40
 
     def __post_init__(self):
         super().__post_init__()

fastvideo/configs/pipelines/base.py

@@ -92,6 +92,9 @@ class PipelineConfig:
     # Compilation
     # enable_torch_compile: bool = False
 
+    # Self-Forcing parameters
+    is_causal: bool = False
+
     @staticmethod
     def add_cli_args(parser: FlexibleArgumentParser,
                      prefix: str = "") -> FlexibleArgumentParser:

fastvideo/configs/pipelines/registry.py

@@ -13,7 +13,7 @@ from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
     Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
     WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
-    SelfForcingWanT2V480PConfig, SelfForcingMoEWanT2V480PConfig)
+    SelfForcingWanT2V480PConfig, SelfForcingMoEWanT2V480PConfig, SelfForcingMoEWanT2V720PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -37,6 +37,8 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
     "rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers": SelfForcingMoEWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingMoEWanT2V720PConfig,
+    "rand0nmr/SFWan2.2-T2V-14B-Diffusers": SelfForcingWanT2V480PConfig,
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
     "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
     "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,

fastvideo/configs/pipelines/wan.py

@@ -164,11 +164,36 @@ class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
     flow_shift: float | None = 5.0
     dmd_denoising_steps: list[int] | None = field(
         default_factory=lambda: [1000, 750, 500, 250])
+    # flow_shift: float | None = 12.0
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000,350,275,200,125])
     warp_denoising_step: bool = True
 
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
 @dataclass
 class SelfForcingMoEWanT2V480PConfig(SelfForcingWanT2V480PConfig):
     boundary_ratio: float | None = 0.875
 
     def __post_init__(self) -> None:
-        self.dit_config.boundary_ratio = self.boundary_ratio
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.dit_config.boundary_ratio = self.boundary_ratio
+
+@dataclass
+class SelfForcingMoEWanT2V720PConfig(SelfForcingMoEWanT2V480PConfig):
+    flow_shift: float | None = 12.0
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000,850,700,550,350,275,200,125])
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000,850,350,275])
+    # flow_shift: float | None = 5.0
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000,850,800,750,570,450,350,250])
+
+    def __post_init__(self) -> None:
+        self.dit_config.boundary_ratio = self.boundary_ratio
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/configs/sample/registry.py

@@ -19,6 +19,7 @@ from fastvideo.configs.sample.wan import (
     WanT2V_1_3B_SamplingParam,
     WanT2V_14B_SamplingParam,
     SelfForcingWanT2V480PConfig,
+    SelfForcingMoEWanT2V480PConfig,
 )
 # isort: on
 from fastvideo.logger import init_logger
@@ -55,7 +56,11 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
 
     # Causal Self-Forcing Wan2.1
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-14B-Diffusers": SelfForcingWanT2V480PConfig,
     "rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+
+    # Causal Self-Forcing Wan2.2 MoE
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingMoEWanT2V480PConfig,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/wan.py

@@ -167,4 +167,17 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
 # =============================================
 @dataclass
 class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
-    pass
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+
+@dataclass
+class SelfForcingMoEWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
+    height: int = 480
+    width: int = 832
+    # height: int = 720
+    # width: int = 1280
+    # num_frames: int = 157
+    num_frames: int = 81
+    fps: int = 16

fastvideo/fastvideo_args.py

@@ -166,6 +166,8 @@ class FastVideoArgs:
 
     # MoE parameters used by Wan2.2
     boundary_ratio: float | None = 0.875
+    override_boundary_ratio: float | None = None
+    use_add_noise_high: bool = False
 
     @property
     def training_mode(self) -> bool:
@@ -411,6 +413,18 @@ class FastVideoArgs:
             default=FastVideoArgs.override_transformer_cls_name,
             help="Override transformer cls name",
         )
+        parser.add_argument(
+            "--override-boundary-ratio",
+            type=float,
+            default=FastVideoArgs.override_boundary_ratio,
+            help="Override boundary ratio for MoE",
+        )
+        parser.add_argument(
+            "--use-add-noise-high",
+            action=StoreBoolean,
+            default=FastVideoArgs.use_add_noise_high,
+            help="Whether to use add_noise_high & pred_noise_to_x_bound for MoE",
+        )
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 

fastvideo/models/loader/component_loader.py

@@ -293,6 +293,8 @@ class TextEncoderLoader(ComponentLoader):
                         mesh=mesh["offload"],
                         fsdp_shard_conditions=model._fsdp_shard_conditions,
                         pin_cpu_memory=fastvideo_args.pin_cpu_memory)
+            else:
+                raise ValueError("text_encoder_cpu_offload must be True")
             # We only enable strict check for non-quantized models
             # that have loaded weights tracking currently.
             # if loaded_weights is not None:
@@ -507,8 +509,9 @@ class SchedulerLoader(ComponentLoader):
         scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
 
         scheduler = scheduler_cls(**config)
-        if fastvideo_args.pipeline_config.flow_shift is not None:
-            scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
+        assert config['shift'] == 12, "shift must be 12"
+        # if fastvideo_args.pipeline_config.flow_shift is not None:
+        #     scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
         if fastvideo_args.pipeline_config.timesteps_scale is not None:
             scheduler.set_timesteps_scale(
                 fastvideo_args.pipeline_config.timesteps_scale)

fastvideo/models/loader/fsdp_load.py

@@ -91,6 +91,7 @@ def maybe_load_fsdp_model(
     )
 
     logger.info("Loading model with default_dtype: %s", default_dtype)
+    logger.info("Loading model with cpu_offload: %s", cpu_offload)
     with set_default_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
 
@@ -222,7 +223,7 @@ def shard_model(
         )
 
     # Finally shard the entire model to account for any stragglers
-    fully_shard(model, **fsdp_kwargs)
+    # fully_shard(model, **fsdp_kwargs)
 
 
 # TODO(PY): device mesh for cfg parallel

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -50,7 +50,8 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
                        stage=CausalDMDDenosingStage(
                            transformer=self.get_module("transformer"),
                            transformer_2=self.get_module("transformer_2", None),
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae")))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae")))

fastvideo/pipelines/composed_pipeline_base.py

@@ -71,7 +71,7 @@ class ComposedPipelineBase(ABC):
 
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
-        fastvideo_args.dit_cpu_offload = False
+        fastvideo_args.dit_cpu_offload = True
         self.modules = self.load_modules(fastvideo_args, loaded_modules)
 
     def set_trainable(self) -> None:

fastvideo/pipelines/pipeline_batch_info.py

@@ -234,6 +234,8 @@ class TrainingBatch:
     # Training outputs
     total_loss: float | None = None
     grad_norm: float | None = None
+    generator_grad_norm: float | None = None
+    critic_grad_norm: float | None = None
 
     # Distillation-specific attributes
     encoder_hidden_states_neg: torch.Tensor | None = None

fastvideo/pipelines/stages/causal_denoising.py

@@ -35,11 +35,12 @@ class CausalDMDDenosingStage(DenoisingStage):
     Denoising stage for causal diffusion.
     """
 
-    def __init__(self, transformer, scheduler, transformer_2=None) -> None:
-        super().__init__(transformer, scheduler, transformer_2)
+    def __init__(self, transformer, scheduler, transformer_2=None, vae=None) -> None:
+        super().__init__(transformer, scheduler, transformer_2, vae)
         # KV and cross-attention cache state (initialized on first forward)
         self.transformer = transformer
         self.transformer_2 = transformer_2
+        self.vae = vae
         # self.kv_cache1: list | None = None
         # self.crossattn_cache: list | None = None
         # Model-dependent constants (aligned with causal_inference.py assumptions)
@@ -80,6 +81,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                                                           dtype=torch.float32)))
             timesteps = scheduler_timesteps[1000 - timesteps]
         timesteps = timesteps.to(get_local_torch_device())
+        logger.info("timesteps: %s", timesteps)
 
         # Image kwargs (kept empty unless caller provides compatible args)
         image_kwargs: dict = {}
@@ -203,39 +205,72 @@ class CausalDMDDenosingStage(DenoisingStage):
         pos_start_base = current_start_frame
 
         # Determine block sizes
-        if not independent_first_frame or (independent_first_frame
-                                           and batch.image_latent is not None):
-            if t % self.num_frames_per_block != 0:
-                raise ValueError(
-                    "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
-                )
-            num_blocks = t // self.num_frames_per_block
-            block_sizes = [self.num_frames_per_block] * num_blocks
-            start_index = 0
-        else:
-            if (t - 1) % self.num_frames_per_block != 0:
-                raise ValueError(
-                    "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
-                )
-            num_blocks = (t - 1) // self.num_frames_per_block
-            block_sizes = [1] + [self.num_frames_per_block] * num_blocks
-            start_index = 0
+        block_sizes = [self.num_frames_per_block] * 7
+        start_index = 0
 
-        if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None:
-            boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps
+        if fastvideo_args.override_boundary_ratio is not None:
+            boundary_timestep = fastvideo_args.override_boundary_ratio * self.scheduler.num_train_timesteps
+            logger.info("Overriding boundary timestep to %s", boundary_timestep)
         else:
-            boundary_timestep = None
-        boundary_timestep = timesteps[2] + 1 # Hardcode for now
+            if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None:
+                boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps
+            else:
+                boundary_timestep = None
 
         high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
-        assert len(high_noise_timesteps) == 2, "only support two high noise timesteps"
+        assert len(high_noise_timesteps) == len(timesteps) // 2, f"There are incorrect number of {len(high_noise_timesteps)} high noise timesteps"
+
+        first_frame_latent = None
+        if batch.pil_image is not None:
+            # Causal video gen directly replaces the first frame of the latent with
+            # the image latent instead of appending along the channel dim
+            assert self.vae is not None, "VAE is not provided for causal video gen task"
+            self.vae = self.vae.to(get_local_torch_device())
+            first_frame_latent = self.vae.encode(batch.pil_image).mean.float()
+            if (hasattr(self.vae, "shift_factor")
+                    and self.vae.shift_factor is not None):
+                if isinstance(self.vae.shift_factor, torch.Tensor):
+                    first_frame_latent -= self.vae.shift_factor.to(first_frame_latent.device, first_frame_latent.dtype)
+                else:
+                    first_frame_latent -= self.vae.shift_factor
+
+            if isinstance(self.vae.scaling_factor, torch.Tensor):
+                first_frame_latent = first_frame_latent * self.vae.scaling_factor.to(first_frame_latent.device, first_frame_latent.dtype)
+            else:
+                first_frame_latent = first_frame_latent * self.vae.scaling_factor
+
+            # latents[:, :, :3, :, :] = first_frame_latent.unsqueeze(2).repeat(1, 1, 3, 1, 1)
+            # start_index += 3
+            # block_sizes.pop(0)
+            
+            if fastvideo_args.vae_cpu_offload:
+                self.vae = self.vae.to("cpu")
+            else:
+                raise ValueError("vae_cpu_offload must be True")
+
+            # if boundary_timestep is not None:
+                # first_frame_latent_bound = self.scheduler.add_noise(
+                #     first_frame_latent,
+                #     torch.randn_like(first_frame_latent),
+                #     boundary_timestep * torch.ones([1], device=first_frame_latent.device, dtype=torch.long))
 
         # DMD loop in causal blocks
         with self.progress_bar(total=len(block_sizes) *
                                len(timesteps)) as progress_bar:
-            for current_num_frames in block_sizes:
+            for block_idx, current_num_frames in enumerate(block_sizes):
                 current_latents = latents[:, :, start_index:start_index +
                                           current_num_frames, :, :]
+
+                if block_idx == 0 and first_frame_latent is not None:
+                    # previous_frame = latents[:, :, start_index - 1:start_index, :, :].clone()
+                    first_frame_mask = torch.ones_like(current_latents)
+                    first_frame_mask[:, :, 0] = 0
+                    current_latents = (1 - first_frame_mask) * first_frame_latent + first_frame_mask * current_latents
+                    # if boundary_timestep is not None:
+                    #     current_latents = first_frame_latent_bound.repeat(1, 1, current_num_frames, 1, 1)
+                    # else:
+                    #     current_latents = first_frame_latent.repeat(1, 1, current_num_frames, 1, 1)
+                
                 # use BTCHW for DMD conversion routines
                 noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
                 video_raw_latent_shape = noise_latents_btchw.shape
@@ -292,9 +327,13 @@ class CausalDMDDenosingStage(DenoisingStage):
                                             forward_batch=batch):
                         # Run transformer; follow DMD stage pattern
                         t_expanded_noise = t_cur * torch.ones(
-                            (latent_model_input.shape[0], 1),
+                            (latent_model_input.shape[0], current_num_frames),
                             device=latent_model_input.device,
                             dtype=torch.long)
+
+                        if block_idx == 0 and first_frame_latent is not None:
+                            t_expanded_noise[:, 0] = 0
+
                         pred_noise_btchw = current_model(
                             latent_model_input,
                             prompt_embeds,
@@ -309,18 +348,18 @@ class CausalDMDDenosingStage(DenoisingStage):
                         ).permute(0, 2, 1, 3, 4)
 
                     # Convert pred noise to pred video with FM Euler scheduler utilities
-                    if boundary_timestep is not None and t_cur >= boundary_timestep:
+                    if boundary_timestep is not None and fastvideo_args.use_add_noise_high and t_cur >= boundary_timestep:
                         pred_video_btchw = pred_noise_to_x_bound(
                             pred_noise=pred_noise_btchw.flatten(0, 1),
                             noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
-                            boundary_timestep=torch.ones_like(t_expand) * boundary_timestep,
+                            timestep=t_expanded_noise.flatten(0, 1),
+                            boundary_timestep=torch.ones_like(t_expanded_noise.flatten(0, 1)) * boundary_timestep,
                             scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2])
                     else:
                         pred_video_btchw = pred_noise_to_pred_video(
                             pred_noise=pred_noise_btchw.flatten(0, 1),
                             noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
+                            timestep=t_expanded_noise.flatten(0, 1),
                             scheduler=self.scheduler).unflatten(
                                 0, pred_noise_btchw.shape[:2])
 
@@ -336,13 +375,13 @@ class CausalDMDDenosingStage(DenoisingStage):
                                 batch.generator, list) else
                                        batch.generator)).to(self.device)
                         noise_btchw = noise
-                        if boundary_timestep is not None and i < len(high_noise_timesteps) - 1:
+                        if boundary_timestep is not None and fastvideo_args.use_add_noise_high and i < len(high_noise_timesteps) - 1:
                             noise_latents_btchw = self.scheduler.add_noise_high(
                                 pred_video_btchw.flatten(0, 1),
                                 noise_btchw.flatten(0, 1),
                                 next_timestep,
                                 torch.ones_like(next_timestep) * boundary_timestep).unflatten(0, pred_video_btchw.shape[:2])
-                        elif boundary_timestep is not None and i == len(high_noise_timesteps) - 1:
+                        elif boundary_timestep is not None and fastvideo_args.use_add_noise_high and i == len(high_noise_timesteps) - 1:
                             noise_latents_btchw = pred_video_btchw
                         else:
                             noise_latents_btchw = self.scheduler.add_noise(
@@ -356,6 +395,11 @@ class CausalDMDDenosingStage(DenoisingStage):
                         current_latents = pred_video_btchw.permute(
                             0, 2, 1, 3, 4)
 
+                    if block_idx == 0 and first_frame_latent is not None:
+                        noise_latents_btchw = (1 - first_frame_mask) * first_frame_latent + first_frame_mask * noise_latents_btchw.permute(0, 2, 1, 3, 4)
+                        noise_latents_btchw = noise_latents_btchw.permute(0, 2, 1, 3, 4)
+                        current_latents = (1 - first_frame_mask) * first_frame_latent + first_frame_mask * current_latents
+
                     if progress_bar is not None:
                         progress_bar.update()
 
@@ -363,6 +407,14 @@ class CausalDMDDenosingStage(DenoisingStage):
                 latents[:, :, start_index:start_index +
                         current_num_frames, :, :] = current_latents
 
+                # Decode the latent
+                # first_frame = block_idx == 0
+                # pixels, self.decode_vae_cache = self.decoding_stage.decode(current_latents, fastvideo_args, *self.decode_vae_cache, first_frame=first_frame)
+                # if batch.output is None:
+                #     batch.output = pixels.cpu().float()
+                # else:
+                #     batch.output = torch.cat([batch.output, pixels.cpu().float()], dim=2)
+
                 # Re-run with context timestep to update KV cache using clean context
                 context_noise = getattr(fastvideo_args.pipeline_config,
                                         "context_noise", 0)
@@ -377,6 +429,21 @@ class CausalDMDDenosingStage(DenoisingStage):
                                         attn_metadata=attn_metadata,
                                         forward_batch=batch):
                     t_expanded_context = t_context.unsqueeze(1)
+
+                    if boundary_timestep is not None:
+                        _ = self.transformer(
+                            context_bcthw,
+                            prompt_embeds,
+                            t_expanded_context,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
+                            current_start=(pos_start_base + start_index) *
+                            self.frame_seq_length,
+                            start_frame=start_index,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        )
+                    
                     _ = current_model(
                         context_bcthw,
                         prompt_embeds,
@@ -389,12 +456,20 @@ class CausalDMDDenosingStage(DenoisingStage):
                         **image_kwargs,
                         **pos_cond_kwargs,
                     )
+
                 start_index += current_num_frames
 
-        # del kv_cache1
-        # del crossattn_cache
-        # gc.collect()
-        # torch.cuda.empty_cache()
+        # remove_indices = torch.arange(3, end=latents.shape[2], step=self.num_frames_per_block)
+        # mask = torch.ones(latents.shape[2], device=latents.device, dtype=torch.bool)
+        # mask[remove_indices] = False
+        # latents = latents[:, :, mask]
+        # latents = latents[:, :, 2:, :, :]
+        # logger.info("latents: %s", latents.shape)
+
+        del kv_cache1
+        del crossattn_cache
+        gc.collect()
+        torch.cuda.empty_cache()
 
         batch.latents = latents
         return batch
@@ -410,6 +485,7 @@ class CausalDMDDenosingStage(DenoisingStage):
             kv_cache_size = self.local_attn_size * self.frame_seq_length
         else:
             kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
+        # kv_cache_size = 1560 * 19
 
         for _ in range(self.num_transformer_blocks):
             kv_cache1.append({

fastvideo/pipelines/stages/decoding.py

@@ -176,7 +176,7 @@ class DecodingStage(PipelineStage):
             self.maybe_free_model_hooks()
 
         if fastvideo_args.vae_cpu_offload:
-            self.vae.to("cpu")
+            self.vae = self.vae.to("cpu")
 
         if torch.backends.mps.is_available():
             del self.vae

fastvideo/pipelines/stages/input_validation.py

@@ -105,27 +105,27 @@ class InputValidationStage(PipelineStage):
                 image = load_image(batch.image_path)
             batch.pil_image = image
 
-        # further processing for ti2v task
-        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+        # further processing for ti2v task or causal video gen task
+        if (fastvideo_args.pipeline_config.ti2v_task or fastvideo_args.pipeline_config.is_causal) and batch.pil_image is not None:
             img = batch.pil_image
             ih, iw = img.height, img.width
             patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
             vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
             dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
-            max_area = 704 * 1280
+            max_area = 480 * 832
             ow, oh = best_output_size(iw, ih, dw, dh, max_area)
 
             scale = max(ow / iw, oh / ih)
             img = img.resize((round(iw * scale), round(ih * scale)),
                              Image.LANCZOS)
-            logger.info("resized img height: %s, img width: %s", img.height,
-                        img.width)
 
             # center-crop
             x1 = (img.width - ow) // 2
             y1 = (img.height - oh) // 2
             img = img.crop((x1, y1, x1 + ow, y1 + oh))
             assert img.width == ow and img.height == oh
+            logger.info("final processed img height: %s, img width: %s",
+                        img.height, img.width)
 
             # to tensor
             img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(

fastvideo/training/checkpointing_utils.py

@@ -21,14 +21,15 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
         state_dict = get_model_state_dict(
             self.model)  # type: ignore[no-any-return]
         # filter out non-trainable parameters
-        param_requires_grad = set([
-            k for k, v in dict(self.model.named_parameters()).items()
-            if v.requires_grad
-        ])
-        state_dict = {
-            k: v
-            for k, v in state_dict.items() if k in param_requires_grad
-        }
+        # param_requires_grad = set([
+        #     k for k, v in dict(self.model.named_parameters()).items()
+        #     if v.requires_grad
+        # ])
+        # state_dict = {
+        #     k: v
+        #     for k, v in state_dict.items() if k in param_requires_grad
+        # }
+        # logger.info(f"state_dict: {state_dict.keys()}")
         return state_dict  # type: ignore
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:

fastvideo/training/distillation_pipeline.py

@@ -39,7 +39,7 @@ from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
     EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
     get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
-    shift_timestep)
+    shift_timestep, unshift_timestep)
 from fastvideo.utils import (is_vsa_available, maybe_download_model,
                              set_random_seed, verify_model_config_and_directory)
 
@@ -69,6 +69,7 @@ class DistillationPipeline(TrainingPipeline):
     video_latent_shape: tuple[int, ...]
     video_latent_shape_sp: tuple[int, ...]
     train_fake_score_transformer_2: bool = False
+    i2v_prompt_image_pairs: list[dict[str, str]] = []
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         raise RuntimeError(
@@ -78,6 +79,10 @@ class DistillationPipeline(TrainingPipeline):
         """Initialize the distillation training pipeline with multiple models."""
         logger.info("Initializing distillation pipeline...")
 
+        json_path = "data/mixkit_i2v_full.json"
+        with open(json_path, 'r') as f:
+            self.i2v_prompt_image_pairs = json.load(f)['data']
+
         super().initialize_training_pipeline(training_args)
 
         self.noise_scheduler = self.get_module("scheduler")
@@ -88,10 +93,14 @@ class DistillationPipeline(TrainingPipeline):
         self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
             shift=self.timestep_shift)
 
-        if self.training_args.boundary_ratio is not None:
-            self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+        if self.training_args.override_boundary_ratio is not None:
+            self.boundary_timestep = self.training_args.override_boundary_ratio * self.noise_scheduler.num_train_timesteps
+            logger.info("Overriding boundary timestep to %s", self.boundary_timestep)
         else:
-            self.boundary_timestep = None
+            if self.training_args.boundary_ratio is not None:
+                self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+            else:
+                self.boundary_timestep = None
 
         if training_args.real_score_model_path:
             logger.info("Loading real score transformer from: %s",
@@ -252,7 +261,6 @@ class DistillationPipeline(TrainingPipeline):
             get_local_torch_device())
         logger.info("Distillation generator model to %s denoising steps: %s",
                     len(self.denoising_step_list), self.denoising_step_list)
-        self.boundary_timestep = self.denoising_step_list[2] + 1 # Hardcode for now
         self.num_train_timestep = self.noise_scheduler.num_train_timesteps
 
         self.min_timestep = int(self.training_args.min_timestep_ratio *
@@ -263,23 +271,23 @@ class DistillationPipeline(TrainingPipeline):
         self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
         self.real_score_guidance_scale_2 = self.training_args.real_score_guidance_scale_2
 
-        self.generator_ema: EMA_FSDP | None = None
-        self.generator_ema_2: EMA_FSDP | None = None
-        if (self.training_args.ema_decay
-                is not None) and (self.training_args.ema_decay > 0.0):
-            self.generator_ema = EMA_FSDP(self.transformer,
-                                          decay=self.training_args.ema_decay)
-            logger.info("Initialized generator EMA with decay=%s",
-                        self.training_args.ema_decay)
+        # self.generator_ema: EMA_FSDP | None = None
+        # self.generator_ema_2: EMA_FSDP | None = None
+        # if (self.training_args.ema_decay
+        #         is not None) and (self.training_args.ema_decay > 0.0):
+        #     self.generator_ema = EMA_FSDP(self.transformer,
+        #                                   decay=self.training_args.ema_decay)
+        #     logger.info("Initialized generator EMA with decay=%s",
+        #                 self.training_args.ema_decay)
 
-            # Initialize EMA for transformer_2 if it exists
-            if self.transformer_2 is not None:
-                self.generator_ema_2 = EMA_FSDP(
-                    self.transformer_2, decay=self.training_args.ema_decay)
-                logger.info("Initialized generator EMA_2 with decay=%s",
-                            self.training_args.ema_decay)
-        else:
-            logger.info("Generator EMA disabled (ema_decay <= 0.0)")
+        #     # Initialize EMA for transformer_2 if it exists
+        #     if self.transformer_2 is not None:
+        #         self.generator_ema_2 = EMA_FSDP(
+        #             self.transformer_2, decay=self.training_args.ema_decay)
+        #         logger.info("Initialized generator EMA_2 with decay=%s",
+        #                     self.training_args.ema_decay)
+        # else:
+        #     logger.info("Generator EMA disabled (ema_decay <= 0.0)")
 
     def load_module_from_path(self, model_path: str, module_type: str,
                               training_args: "TrainingArgs"):
@@ -358,9 +366,11 @@ class DistillationPipeline(TrainingPipeline):
         """Prepare training environment for distillation."""
         self.transformer.requires_grad_(True)
         self.transformer.train()
+        self.transformer.scale_shift_table.requires_grad_(False)
         if self.transformer_2 is not None:
             self.transformer_2.requires_grad_(True)
             self.transformer_2.train()
+            self.transformer_2.scale_shift_table.requires_grad_(False)
         self.fake_score_transformer.requires_grad_(True)
         self.fake_score_transformer.train()
         if self.fake_score_transformer_2 is not None:
@@ -754,37 +764,49 @@ class DistillationPipeline(TrainingPipeline):
         """Compute DMD (Diffusion Model Distillation) loss."""
         original_latent = generator_pred_video
         clean_original_latent = clean_generator_pred_video
+        boundary_timestep_pre_shift = int(unshift_timestep(self.boundary_timestep, self.timestep_shift))
         with torch.no_grad():
             # Detect whether we're doing MoE DMD
-            # if self.boundary_timestep is not None:
-            #     if exit_timestep < self.boundary_timestep:
-            #         start_timestep = 0
-            #         # end_timestep = int(self.boundary_timestep)
-            #         end_timestep = self.num_train_timestep
-            #     else:
-            #         # raise ValueError("Exit timestep is greater than boundary timestep")
-            #         start_timestep = int(self.boundary_timestep)
-            #         end_timestep = self.num_train_timestep
-            # else:
-            start_timestep = 0
-            end_timestep = self.num_train_timestep
+            num_high_noise_timesteps = torch.sum(self.noise_scheduler.timesteps >= self.boundary_timestep).item()
+            if self.boundary_timestep is not None:
+                if exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is not None:
+                    train_critic_2 = torch.rand(1, generator=self.noise_random_generator).item() >= (num_high_noise_timesteps / self.num_train_timestep)
+                    decision = torch.tensor(1.0 if train_critic_2 else 0.0,
+                                    device=self.device)
+                    world_group = get_world_group()
+                    if world_group.world_size > 1:
+                        world_group.broadcast(decision, src=0)
+                    train_critic_2 = decision.item() == 1.0
+
+                    if train_critic_2:
+                        start_timestep = 0
+                        end_timestep = boundary_timestep_pre_shift
+                    else:
+                        start_timestep = boundary_timestep_pre_shift + 2
+                        end_timestep = self.num_train_timestep
+                elif exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is None:
+                    start_timestep = 0
+                    end_timestep = self.num_train_timestep
+                else:
+                    start_timestep = boundary_timestep_pre_shift + 2
+                    end_timestep = self.num_train_timestep
+            else:
+                start_timestep = 0
+                end_timestep = self.num_train_timestep
 
             timestep = torch.randint(start_timestep,
                                      end_timestep, [1],
                                      device=self.device,
                                      dtype=torch.long)
-            world_group = get_world_group()
-            if world_group.world_size > 1:
-                world_group.broadcast(timestep, src=0)
 
             timestep = shift_timestep(
                 timestep,
-                self.timestep_shift,  # type: ignore
-                min_timestep=start_timestep,
-                max_timestep=end_timestep)
+                self.timestep_shift)  # type: ignore
 
             timestep = timestep.clamp(self.min_timestep, self.max_timestep)
 
+            # logger.info(f"_dmd_forward: rank {self.global_rank} timestep: {timestep}", local_main_process_only=False)
+
             noise = torch.randn(self.video_latent_shape,
                                 device=self.device,
                                 dtype=generator_pred_video.dtype)
@@ -794,18 +816,20 @@ class DistillationPipeline(TrainingPipeline):
                                   n=self.sp_world_size).contiguous()
                 noise = noise[:, self.rank_in_sp_group, :, :, :, :]
 
-            if self.boundary_timestep is not None and self.transformer_2 is not None and exit_timestep >= self.boundary_timestep:
-                if timestep < self.boundary_timestep:
-                    noisy_latent = self.noise_scheduler.add_noise(
-                        clean_generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
-                        timestep).detach().unflatten(0,
-                                                    (1, clean_original_latent.shape[1]))
-                else:
-                    noisy_latent = self.noise_scheduler.add_noise_high(
-                        generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
-                        timestep,
-                        self.boundary_timestep * torch.ones_like(timestep)).detach().unflatten(0,
-                                                                                            (1, generator_pred_video.shape[1]))
+            moe_use_add_noise_high = self.boundary_timestep is not None and self.training_args.use_add_noise_high and exit_timestep >= self.boundary_timestep
+
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     noisy_latent = self.noise_scheduler.add_noise(
+                #         clean_generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
+                #         timestep).detach().unflatten(0,
+                #                                     (1, clean_original_latent.shape[1]))
+                # else:
+                noisy_latent = self.noise_scheduler.add_noise_high(
+                    generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
+                    timestep,
+                    self.boundary_timestep * torch.ones_like(timestep)).detach().unflatten(0,
+                                                                                        (1, generator_pred_video.shape[1]))
             else:
                 noisy_latent = self.noise_scheduler.add_noise(
                     generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
@@ -821,20 +845,26 @@ class DistillationPipeline(TrainingPipeline):
             fake_score_pred_noise = current_fake_score_transformer(
                 **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
 
-            if self.boundary_timestep is not None and self.transformer_2 is not None and exit_timestep >= self.boundary_timestep:
-                if timestep < self.boundary_timestep:
-                    faker_score_pred_video = pred_noise_to_pred_video(
-                        pred_noise=fake_score_pred_noise.flatten(0, 1),
-                        noise_input_latent=noisy_latent.flatten(0, 1),
-                        timestep=timestep,
-                        scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
-                else:
-                    faker_score_pred_video = pred_noise_to_x_bound(
-                        pred_noise=fake_score_pred_noise.flatten(0, 1),
-                        noise_input_latent=noisy_latent.flatten(0, 1),
-                        timestep=timestep,
-                        boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
-                        scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     raise ValueError("Timestep is less than boundary timestep")
+                #     faker_score_pred_video = pred_noise_to_pred_video(
+                #         pred_noise=fake_score_pred_noise.flatten(0, 1),
+                #         noise_input_latent=noisy_latent.flatten(0, 1),
+                #         timestep=timestep,
+                #         scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+                # else:
+                faker_score_pred_video = pred_noise_to_x_bound(
+                    pred_noise=fake_score_pred_noise.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+                clean_faker_score_pred_video = pred_noise_to_pred_video(
+                    pred_noise=fake_score_pred_noise.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
             else:
                 faker_score_pred_video = pred_noise_to_pred_video(
                     pred_noise=fake_score_pred_noise.flatten(0, 1),
@@ -852,20 +882,26 @@ class DistillationPipeline(TrainingPipeline):
             real_score_pred_noise_cond = current_real_score_transformer(
                 **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
 
-            if self.boundary_timestep is not None and self.transformer_2 is not None and exit_timestep >= self.boundary_timestep:
-                if timestep < self.boundary_timestep:
-                    pred_real_video_cond = pred_noise_to_pred_video(
-                        pred_noise=real_score_pred_noise_cond.flatten(0, 1),
-                        noise_input_latent=noisy_latent.flatten(0, 1),
-                        timestep=timestep,
-                        scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_cond.shape[:2])
-                else:
-                    pred_real_video_cond = pred_noise_to_x_bound(
-                        pred_noise=real_score_pred_noise_cond.flatten(0, 1),
-                        noise_input_latent=noisy_latent.flatten(0, 1),
-                        timestep=timestep,
-                        boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
-                        scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_cond.shape[:2])
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     pred_real_video_cond = pred_noise_to_pred_video(
+                #         pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                #         noise_input_latent=noisy_latent.flatten(0, 1),
+                #         timestep=timestep,
+                #         scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_cond.shape[:2])
+                # else:
+                pred_real_video_cond = pred_noise_to_x_bound(
+                    pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_cond.shape[:2])
+                clean_pred_real_video_cond = pred_noise_to_pred_video(
+                    pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, real_score_pred_noise_cond.shape[:2])
             else:
                 pred_real_video_cond = pred_noise_to_pred_video(
                     pred_noise=real_score_pred_noise_cond.flatten(0, 1),
@@ -882,20 +918,26 @@ class DistillationPipeline(TrainingPipeline):
             real_score_pred_noise_uncond = current_real_score_transformer(
                 **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
 
-            if self.boundary_timestep is not None and self.transformer_2 is not None and exit_timestep >= self.boundary_timestep:
-                if timestep < self.boundary_timestep:
-                    pred_real_video_uncond = pred_noise_to_pred_video(
-                        pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
-                        noise_input_latent=noisy_latent.flatten(0, 1),
-                        timestep=timestep,
-                        scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_uncond.shape[:2])
-                else:
-                    pred_real_video_uncond = pred_noise_to_x_bound(
-                        pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
-                        noise_input_latent=noisy_latent.flatten(0, 1),
-                        timestep=timestep,
-                        boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
-                        scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_uncond.shape[:2])
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     pred_real_video_uncond = pred_noise_to_pred_video(
+                #         pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                #         noise_input_latent=noisy_latent.flatten(0, 1),
+                #         timestep=timestep,
+                #         scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_uncond.shape[:2])
+                # else:
+                pred_real_video_uncond = pred_noise_to_x_bound(
+                    pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_uncond.shape[:2])
+                clean_pred_real_video_uncond = pred_noise_to_pred_video(
+                    pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, real_score_pred_noise_uncond.shape[:2])
             else:
                 pred_real_video_uncond = pred_noise_to_pred_video(
                     pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
@@ -904,6 +946,10 @@ class DistillationPipeline(TrainingPipeline):
                     scheduler=self.noise_scheduler).unflatten(
                         0, real_score_pred_noise_uncond.shape[:2])
 
+            if moe_use_add_noise_high:
+                clean_real_score_pred_video = clean_pred_real_video_cond + (
+                    clean_pred_real_video_cond -
+                    clean_pred_real_video_uncond) * self._get_real_score_guidance_scale(timestep)
             real_score_pred_video = pred_real_video_cond + (
                 pred_real_video_cond -
                 pred_real_video_uncond) * self._get_real_score_guidance_scale(timestep)
@@ -916,18 +962,32 @@ class DistillationPipeline(TrainingPipeline):
             original_latent.float(),
             (original_latent.float() - grad.float()).detach())
 
-        training_batch.dmd_latent_vis_dict.update({
-            "training_batch_dmd_fwd_clean_latent":
-            training_batch.latents,
-            "generator_pred_video":
-            original_latent.detach(),
-            "real_score_pred_video":
-            real_score_pred_video.detach(),
-            "faker_score_pred_video":
-            faker_score_pred_video.detach(),
-            "dmd_timestep":
-            timestep.detach(),
-        })
+        if moe_use_add_noise_high:
+            training_batch.dmd_latent_vis_dict.update({
+                "training_batch_dmd_fwd_clean_latent":
+                training_batch.latents,
+                "generator_pred_video":
+                clean_original_latent.detach(),
+                "real_score_pred_video":
+                clean_real_score_pred_video.detach(),
+                "faker_score_pred_video":
+                clean_faker_score_pred_video.detach(),
+                "dmd_timestep":
+                timestep.detach(),
+            })
+        else:
+            training_batch.dmd_latent_vis_dict.update({
+                "training_batch_dmd_fwd_clean_latent":
+                training_batch.latents,
+                "generator_pred_video":
+                clean_original_latent.detach(),
+                "real_score_pred_video":
+                real_score_pred_video.detach(),
+                "faker_score_pred_video":
+                faker_score_pred_video.detach(),
+                "dmd_timestep":
+                timestep.detach(),
+            })
 
         return dmd_loss
 
@@ -943,38 +1003,50 @@ class DistillationPipeline(TrainingPipeline):
             else:
                 generator_pred_video = self._generator_forward(training_batch)
 
+        num_high_noise_timesteps = torch.sum(self.noise_scheduler.timesteps >= self.boundary_timestep).item()
         # Detect whether we're doing MoE DMD
         # If yes, then we need to use the appropriate timestep range
-        # if self.boundary_timestep is not None:
-        #     if exit_timestep < self.boundary_timestep:
-        #         start_timestep = 0
-        #         end_timestep = self.num_train_timestep
-        #         # end_timestep = int(self.boundary_timestep)
-        #     else:
-        #         # raise ValueError("Exit timestep is greater than boundary timestep")
-        #         start_timestep = int(self.boundary_timestep)
-        #         end_timestep = self.num_train_timestep
-        # else:
-        start_timestep = 0
-        end_timestep = self.num_train_timestep
+        boundary_timestep_pre_shift = int(unshift_timestep(self.boundary_timestep, self.timestep_shift))
+        if self.boundary_timestep is not None:
+            if exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is not None:
+                train_critic_2 = torch.rand(1, generator=self.noise_random_generator).item() >= (num_high_noise_timesteps / self.num_train_timestep)
+                decision = torch.tensor(1.0 if train_critic_2 else 0.0,
+                                device=self.device)
+                world_group = get_world_group()
+                if world_group.world_size > 1:
+                    world_group.broadcast(decision, src=0)
+                train_critic_2 = decision.item() == 1.0
+
+                if train_critic_2:
+                    start_timestep = 0
+                    end_timestep = boundary_timestep_pre_shift
+                else:
+                    start_timestep = boundary_timestep_pre_shift + 2
+                    end_timestep = self.num_train_timestep
+            elif exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is None:
+                start_timestep = 0
+                end_timestep = self.num_train_timestep
+            else:
+                start_timestep = boundary_timestep_pre_shift + 2
+                end_timestep = self.num_train_timestep
+        else:
+            start_timestep = 0
+            end_timestep = self.num_train_timestep
 
         fake_score_timestep = torch.randint(start_timestep,
                                             end_timestep, [1],
                                             device=self.device,
                                             dtype=torch.long)
-        world_group = get_world_group()
-        if world_group.world_size > 1:
-            world_group.broadcast(fake_score_timestep, src=0)
 
         fake_score_timestep = shift_timestep(
             fake_score_timestep,
-            self.timestep_shift,  # type: ignore
-            min_timestep=start_timestep,
-            max_timestep=end_timestep)
+            self.timestep_shift)  # type: ignore
 
         fake_score_timestep = fake_score_timestep.clamp(self.min_timestep,
                                                         self.max_timestep)
 
+        # logger.info(f"faker_score_forward: rank {self.global_rank} fake_score_timestep: {fake_score_timestep}", local_main_process_only=False)
+
         fake_score_noise = torch.randn(self.video_latent_shape,
                                        device=self.device,
                                        dtype=generator_pred_video.dtype)
@@ -984,19 +1056,21 @@ class DistillationPipeline(TrainingPipeline):
                                          n=self.sp_world_size).contiguous()
             fake_score_noise = fake_score_noise[:, self.
                                                 rank_in_sp_group, :, :, :, :]
+        
+        moe_use_add_noise_high = self.boundary_timestep is not None and self.training_args.use_add_noise_high and exit_timestep >= self.boundary_timestep
 
-        if self.boundary_timestep is not None and self.transformer_2 is not None and exit_timestep >= self.boundary_timestep:
-            if fake_score_timestep < self.boundary_timestep:
-                noisy_generator_pred_video = self.noise_scheduler.add_noise(
-                    clean_generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
-                    fake_score_timestep).unflatten(0,
-                                                (1, clean_generator_pred_video.shape[1]))
-            else:
-                noisy_generator_pred_video = self.noise_scheduler.add_noise_high(
-                    generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
-                    fake_score_timestep,
-                    self.boundary_timestep * torch.ones_like(fake_score_timestep)).unflatten(0,
-                                                                                        (1, generator_pred_video.shape[1]))
+        if moe_use_add_noise_high:
+            # if fake_score_timestep < self.boundary_timestep:
+            #     noisy_generator_pred_video = self.noise_scheduler.add_noise(
+            #         clean_generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
+            #         fake_score_timestep).unflatten(0,
+            #                                     (1, clean_generator_pred_video.shape[1]))
+            # else:
+            noisy_generator_pred_video = self.noise_scheduler.add_noise_high(
+                generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
+                fake_score_timestep,
+                self.boundary_timestep * torch.ones_like(fake_score_timestep)).unflatten(0,
+                                                                                    (1, generator_pred_video.shape[1]))
         else:
             noisy_generator_pred_video = self.noise_scheduler.add_noise(
                 generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
@@ -1014,45 +1088,45 @@ class DistillationPipeline(TrainingPipeline):
             fake_score_pred_noise = current_fake_score_transformer(
                 **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)
 
-        if self.boundary_timestep is not None and self.transformer_2 is not None and exit_timestep >= self.boundary_timestep:
-            if fake_score_timestep < self.boundary_timestep:
-                pred_video = pred_noise_to_pred_video(
-                    pred_noise=fake_score_pred_noise.flatten(0, 1),
-                    noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
-                    timestep=fake_score_timestep,
-                    scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
-                flow_matching_loss = torch.mean((pred_video - clean_generator_pred_video)**2)
-            else:
-                pred_video = pred_noise_to_x_bound(
-                    pred_noise=fake_score_pred_noise.flatten(0, 1),
-                    noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
-                    timestep=fake_score_timestep,
-                    boundary_timestep=torch.ones_like(fake_score_timestep) * self.boundary_timestep,
-                    scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
-                flow_matching_loss = torch.mean((pred_video - generator_pred_video)**2)
-        else:
+        if moe_use_add_noise_high:
+            # if fake_score_timestep < self.boundary_timestep:
+            #     pred_video = pred_noise_to_pred_video(
+            #         pred_noise=fake_score_pred_noise.flatten(0, 1),
+            #         noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
+            #         timestep=fake_score_timestep,
+            #         scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            #     flow_matching_loss = torch.mean((pred_video - clean_generator_pred_video)**2)
+            # else:
+            pred_video = pred_noise_to_x_bound(
+                pred_noise=fake_score_pred_noise.flatten(0, 1),
+                noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
+                timestep=fake_score_timestep,
+                boundary_timestep=torch.ones_like(fake_score_timestep) * self.boundary_timestep,
+                scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            flow_matching_loss = torch.mean((pred_video - generator_pred_video)**2)
+        elif self.boundary_timestep is not None and self.training_args.use_add_noise_high and exit_timestep < self.boundary_timestep:
             pred_video = pred_noise_to_pred_video(
                 pred_noise=fake_score_pred_noise.flatten(0, 1),
                 noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
                 timestep=fake_score_timestep,
                 scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
             flow_matching_loss = torch.mean((pred_video - generator_pred_video)**2)
-
-        # This is mathematically equivalent to pred_video - generator_pred_video for the low noise case
-        # target = fake_score_noise - generator_pred_video
-        # flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)
+        else:
+            # This is mathematically equivalent to pred_video - generator_pred_video for the low noise case
+            target = fake_score_noise - generator_pred_video
+            flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)
 
         training_batch.fake_score_latent_vis_dict = {
             "training_batch_fakerscore_fwd_clean_latent":
             training_batch.latents,
-            "generator_pred_video": generator_pred_video,
+            "generator_pred_video": clean_generator_pred_video,
             "fake_score_timestep": fake_score_timestep,
         }
 
         return training_batch, flow_matching_loss
 
     def _clip_model_grad_norm_(self, training_batch: TrainingBatch,
-                               transformer) -> TrainingBatch:
+                               transformer) -> float:
 
         max_grad_norm = self.training_args.max_grad_norm
 
@@ -1068,8 +1142,7 @@ class DistillationPipeline(TrainingPipeline):
             grad_norm = grad_norm.item() if grad_norm is not None else 0.0
         else:
             grad_norm = 0.0
-        training_batch.grad_norm = grad_norm
-        return training_batch
+        return grad_norm
 
     def _prepare_dit_inputs(self,
                             training_batch: TrainingBatch) -> TrainingBatch:
@@ -1334,18 +1407,17 @@ class DistillationPipeline(TrainingPipeline):
         else:
             logger.info("  Generator EMA disabled")
 
-        if self.generator_ema is not None:
-            logger.info("  Generator EMA enabled with decay: %s",
+        if getattr(self, 'generator_ema_2', None) is not None:
+            logger.info("  Generator EMA_2 enabled with decay: %s",
                         self.training_args.ema_decay)
-            logger.info("  Generator EMA start step: %s",
+            logger.info("  Generator EMA_2 start step: %s",
                         self.training_args.ema_start_step)
         else:
-            logger.info("  Generator EMA disabled")
+            logger.info("  Generator EMA_2 disabled")
 
     @torch.no_grad()
-    def _log_validation(self, transformer, training_args, global_step) -> None:
+    def _log_validation(self, training_args, global_step) -> None:
         training_args.inference_mode = True
-        training_args.dit_cpu_offload = True
         if not training_args.log_validation:
             return
         if self.validation_pipeline is None:
@@ -1372,7 +1444,7 @@ class DistillationPipeline(TrainingPipeline):
                                            num_workers=0)
 
         # Set both transformers to eval mode
-        transformer.eval()
+        self.transformer.eval()
         if hasattr(self, 'transformer_2') and self.transformer_2 is not None:
             self.transformer_2.eval()
 
@@ -1383,23 +1455,19 @@ class DistillationPipeline(TrainingPipeline):
         ema_2_context = None
 
         if use_ema_for_validation:
-            logger.info("Using EMA model for validation")
             # Use self.transformer for consistency (the passed transformer should be self.transformer anyway)
-            validation_transformer = self.transformer
-            if self.generator_ema is not None:
-                ema_context = self.generator_ema.apply_to_model(
-                    validation_transformer)
+            # if self.generator_ema is not None:
+            ema_context = self.generator_ema.apply_to_model(
+                self.transformer)
+            logger.info("Using EMA model for validation")
 
             # Handle transformer_2 EMA if available
             if hasattr(
                     self, 'transformer_2'
-            ) and self.transformer_2 is not None and self.generator_ema_2 is not None:
+            ) and self.transformer_2 is not None:
                 ema_2_context = self.generator_ema_2.apply_to_model(
                     self.transformer_2)
                 logger.info("Using EMA_2 model for transformer_2 validation")
-        else:
-            # Use self.transformer for consistency, but the passed transformer should be the same
-            validation_transformer = self.transformer
 
         validation_steps = training_args.validation_sampling_steps.split(",")
         validation_steps = [int(step) for step in validation_steps]
@@ -1409,13 +1477,64 @@ class DistillationPipeline(TrainingPipeline):
         num_sp_groups = world_group.world_size // self.sp_group.world_size
         # Process each validation prompt for each validation step
         for num_inference_steps in validation_steps:
-            logger.info("rank: %s: num_inference_steps: %s",
-                        self.global_rank,
-                        num_inference_steps,
-                        local_main_process_only=False)
+            # logger.info("rank: %s: num_inference_steps: %s",
+            #             self.global_rank,
+            #             num_inference_steps,
+            #             local_main_process_only=False)
             step_videos: list[np.ndarray] = []
             step_captions: list[str] = []
 
+            def run_i2v_validation_with_ema(
+                steps: int) -> tuple[list[np.ndarray], list[str]]:
+                videos: list[np.ndarray] = []
+                captions: list[str] = []
+                for iter in range(len(self.i2v_prompt_image_pairs) // world_group.world_size):
+                    validation_batch = self.i2v_prompt_image_pairs[iter * world_group.world_size + self.global_rank]
+                    batch = self._prepare_validation_batch(
+                        sampling_param, training_args, validation_batch, steps)
+
+                    negative_prompt = batch.negative_prompt
+                    batch_negative = ForwardBatch(
+                        data_type="video",
+                        prompt=negative_prompt,
+                        prompt_embeds=[],
+                        prompt_attention_mask=[],
+                    )
+                    result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                        batch_negative, training_args)
+                    self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                        0], result_batch.prompt_attention_mask[0]
+
+                    logger.info(
+                        "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                        self.global_rank,
+                        self.rank_in_sp_group,
+                        batch.prompt,
+                        local_main_process_only=False)
+
+                    assert batch.prompt is not None and isinstance(
+                        batch.prompt, str)
+                    captions.append(batch.prompt)
+
+                    # Run validation inference
+                    with torch.no_grad():
+                        output_batch = self.validation_pipeline.forward(
+                            batch, training_args)
+                    samples = output_batch.output
+                    if self.rank_in_sp_group != 0:
+                        return videos, captions
+
+                    # Process outputs
+                    video = rearrange(samples, "b c t h w -> t b c h w")
+                    frames = []
+                    for x in video:
+                        x = torchvision.utils.make_grid(x, nrow=6)
+                        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                        frames.append((x * 255).numpy().astype(np.uint8))
+                    videos.append(frames)
+
+                return videos, captions
+
             # Helper function to run validation with optional EMA contexts
             def run_validation_with_ema(
                     steps: int) -> tuple[list[np.ndarray], list[str]]:
@@ -1484,6 +1603,22 @@ class DistillationPipeline(TrainingPipeline):
                 step_videos, step_captions = run_validation_with_ema(
                     num_inference_steps)
 
+            if ema_context is not None and ema_2_context is not None:
+                with ema_context, ema_2_context:
+                    i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                        num_inference_steps)
+            elif ema_context is not None:
+                with ema_context:
+                    i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                        num_inference_steps)
+            elif ema_2_context is not None:
+                with ema_2_context:
+                    i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                        num_inference_steps)
+            else:
+                i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                    num_inference_steps)
+
             # Log validation results for this step
             world_group = get_world_group()
             num_sp_groups = world_group.world_size // self.sp_group.world_size
@@ -1528,8 +1663,47 @@ class DistillationPipeline(TrainingPipeline):
                     world_group.send_object(step_videos, dst=0)
                     world_group.send_object(step_captions, dst=0)
 
+                # sp_group leaders wait for global rank 0 to finish logging
+                world_group.barrier()
+
+                if self.global_rank == 0:
+                    all_i2v_videos = i2v_step_videos
+                    all_i2v_captions = i2v_step_captions
+
+                    # Receive from other sp_group leaders
+                    for sp_group_idx in range(1, num_sp_groups):
+                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                        recv_videos = world_group.recv_object(src=src_rank)
+                        recv_captions = world_group.recv_object(src=src_rank)
+                        all_i2v_videos.extend(recv_videos)
+                        all_i2v_captions.extend(recv_captions)
+
+                    video_filenames = []
+                    for i, (video, caption) in enumerate(
+                            zip(all_i2v_videos, all_i2v_captions, strict=True)):
+                        os.makedirs(training_args.output_dir, exist_ok=True)
+                        filename = os.path.join(
+                            training_args.output_dir,
+                            f"validation_step_{global_step}_inference_steps_{num_inference_steps}_i2v_video_{i}.mp4"
+                        )
+                        imageio.mimsave(filename, video, fps=sampling_param.fps)
+                        video_filenames.append(filename)
+
+                    logs = {
+                        f"i2v_validation_videos_{num_inference_steps}_steps": [
+                            wandb.Video(filename, caption=caption)
+                            for filename, caption in zip(
+                                video_filenames, all_i2v_captions, strict=True)
+                        ]
+                    }
+                    wandb.log(logs, step=global_step)
+                else:
+                    world_group.send_object(i2v_step_videos, dst=0)
+                    world_group.send_object(i2v_step_captions, dst=0)
+
+
         # Re-enable gradients for training - set both transformers back to train mode
-        transformer.train()
+        self.transformer.train()
         if hasattr(self, 'transformer_2') and self.transformer_2 is not None:
             self.transformer_2.train()
         gc.collect()
@@ -1854,6 +2028,8 @@ class DistillationPipeline(TrainingPipeline):
                 if self.training_args.use_ema and self.is_ema_ready():
                     self.save_ema_weights(self.training_args.output_dir, step)
 
+                self.sp_group.barrier()
+
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                 if self.training_args.log_visualization:
                     self.visualize_intermediate_latents(training_batch,

fastvideo/training/self_forcing_distillation_pipeline.py

@@ -27,6 +27,7 @@ from fastvideo.training.distillation_pipeline import DistillationPipeline
 from fastvideo.training.training_utils import (EMA_FSDP,
                                                save_distillation_checkpoint)
 from fastvideo.utils import is_vsa_available, set_random_seed
+from copy import deepcopy
 
 logger = init_logger(__name__)
 
@@ -48,10 +49,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         """Initialize the self-forcing training pipeline."""
         logger.info("Initializing self-forcing distillation pipeline...")
 
+        # self.generator_ema: EMA_FSDP | None = None
+        # self.generator_ema_2: EMA_FSDP | None = None
+        
+        super().initialize_training_pipeline(training_args)
+
         self.generator_ema: EMA_FSDP | None = None
         self.generator_ema_2: EMA_FSDP | None = None
 
-        super().initialize_training_pipeline(training_args)
         try:
             logger.info("RANK: %s, entered initialize_training_pipeline",
                         self.global_rank,
@@ -61,7 +66,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
         self.noise_scheduler = SelfForcingFlowMatchScheduler(
             num_inference_steps=1000,
-            shift=5.0,
+            shift=self.timestep_shift,
             sigma_min=0.0,
             extra_one_step=True,
             training=True)
@@ -89,13 +94,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
     def generate_and_sync_list(self, num_blocks: int, num_denoising_steps: int,
                                device: torch.device) -> list[int]:
         """Generate and synchronize random exit flags across distributed processes."""
-        logger.info(
-            "RANK: %s, enter generate_and_sync_list blocks=%s steps=%s device=%s",
-            self.global_rank,
-            num_blocks,
-            num_denoising_steps,
-            str(device),
-            local_main_process_only=False)
+        # logger.info(
+        #     "RANK: %s, enter generate_and_sync_list blocks=%s steps=%s device=%s",
+        #     self.global_rank,
+        #     num_blocks,
+        #     num_denoising_steps,
+        #     str(device),
+        #     local_main_process_only=False)
         rank = dist.get_rank() if dist.is_initialized() else 0
 
         if rank == 0:
@@ -113,12 +118,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             dist.broadcast(indices,
                            src=0)  # Broadcast the random indices to all ranks
         flags = indices.tolist()
-        logger.info(
-            "RANK: %s, exit generate_and_sync_list flags_len=%s first=%s",
-            self.global_rank,
-            len(flags),
-            flags[0] if len(flags) > 0 else None,
-            local_main_process_only=False)
+        # logger.info(
+        #     "RANK: %s, exit generate_and_sync_list flags_len=%s first=%s",
+        #     self.global_rank,
+        #     len(flags),
+        #     flags[0] if len(flags) > 0 else None,
+        #     local_main_process_only=False)
         return flags
 
     def generator_loss(
@@ -128,6 +133,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         Compute generator loss using DMD-style approach.
         The generator tries to fool the critic (fake_score_transformer).
         """
+        # Check that all params in self.transformer require grad is True
+        # for param in self.transformer.parameters():
+        #     assert param.requires_grad, "Param %s in transformer requires grad is False" % param.name
+        # if self.transformer_2 is not None:
+        #     for param in self.transformer_2.parameters():
+        #         assert param.requires_grad, "Param %s in transformer_2 requires grad is False" % param.name
         exit_flags = None
         # If the generator is MoE, randomly sample an exit timestep, and turn on training for one of the dits
         if self.boundary_timestep is not None and self.transformer_2 is not None:
@@ -139,20 +150,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             # logger.info("Exit timestep in generator_loss(): %s", exit_timestep)
             if exit_timestep < self.boundary_timestep:
                 self.train_transformer_2 = True
-                self._enable_training(self.transformer_2, self.optimizer_2)
-                self._disable_training(self.transformer, self.optimizer)
             else:
                 self.train_transformer_2 = False
-                self._enable_training(self.transformer, self.optimizer)
-                self._disable_training(self.transformer_2, self.optimizer_2)
         else: # Default to normal single-dit generator
             self.train_transformer_2 = False
-            self._enable_training(self.transformer, self.optimizer)
-
-        # Turns off training for fake score transformers
-        self._disable_training(self.fake_score_transformer, self.fake_score_optimizer)
-        if self.fake_score_transformer_2 is not None:
-            self._disable_training(self.fake_score_transformer_2, self.fake_score_optimizer_2)
 
         with set_forward_context(
                 current_timestep=training_batch.timesteps,
@@ -182,15 +183,11 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         Compute critic loss using flow matching between noise and generator output.
         The critic learns to predict the flow from noise to the generator's output.
         """
-        # Turns off training for all generators
-        self._disable_training(self.transformer, self.optimizer)
-        if self.transformer_2 is not None:
-            self._disable_training(self.transformer_2, self.optimizer_2)
-
-        # Turns on training for fake score transformers
-        self._enable_training(self.fake_score_transformer, self.fake_score_optimizer)
-        if self.fake_score_transformer_2 is not None:
-            self._enable_training(self.fake_score_transformer_2, self.fake_score_optimizer_2)
+        # for param in self.fake_score_transformer.parameters():
+        #     assert param.requires_grad, "Param %s in fake_score_transformer requires grad is False" % param.name
+        # if self.fake_score_transformer_2 is not None:
+        #     for param in self.fake_score_transformer_2.parameters():
+        #         assert param.requires_grad, "Param %s in fake_score_transformer_2 requires grad is False" % param.name
 
         updated_batch, flow_matching_loss = self.faker_score_forward(
             training_batch)
@@ -213,29 +210,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         batch_size = latents.shape[0]
         initial_latent = getattr(training_batch, 'image_latent', None)
 
-        # Dynamic frame generation logic (adapted from _run_generator)
-        num_training_frames = getattr(self.training_args, 'num_latent_t', 21)
-
-        # During training, the number of generated frames should be uniformly sampled from
-        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
-        min_num_frames = 20 if self.independent_first_frame else 21
-        max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames
-        assert max_num_frames % self.num_frame_per_block == 0
-        assert min_num_frames % self.num_frame_per_block == 0
-        max_num_blocks = max_num_frames // self.num_frame_per_block
-        min_num_blocks = min_num_frames // self.num_frame_per_block
-
-        # Sample number of blocks and sync across processes
-        num_generated_blocks = torch.randint(min_num_blocks,
-                                             max_num_blocks + 1, (1, ),
-                                             device=self.device)
-        if dist.is_initialized():
-            dist.broadcast(num_generated_blocks, src=0)
-        num_generated_blocks = num_generated_blocks.item()
-        num_generated_frames = num_generated_blocks * self.num_frame_per_block
-        if self.independent_first_frame and initial_latent is None:
-            num_generated_frames += 1
-            min_num_frames += 1
+        num_generated_frames = getattr(self.training_args, 'num_latent_t', 21)
 
         # Create noise with dynamic shape
         if initial_latent is not None:
@@ -258,13 +233,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         batch_size, num_frames, num_channels, height, width = noise.shape
 
         # Block size calculation
-        if not self.independent_first_frame or (self.independent_first_frame
-                                                and initial_latent is not None):
-            assert num_frames % self.num_frame_per_block == 0
-            num_blocks = num_frames // self.num_frame_per_block
-        else:
-            assert (num_frames - 1) % self.num_frame_per_block == 0
-            num_blocks = (num_frames - 1) // self.num_frame_per_block
 
         num_input_frames = initial_latent.shape[
             1] if initial_latent is not None else 0
@@ -278,14 +246,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             device=noise.device,
             dtype=noise.dtype)
 
-        def get_model_device(model):
-            if model is None:
-                return "None"
-            try:
-                return next(model.parameters()).device
-            except (StopIteration, AttributeError):
-                return "Unknown"
-
         # Step 1: Initialize KV cache to all zeros
         cache_frames = num_generated_frames + num_input_frames
         kv_cache1, crossattn_cache = self._initialize_simulation_caches(
@@ -319,7 +279,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             current_start_frame += 1
 
         # Step 3: Temporal denoising loop
-        all_num_frames = [self.num_frame_per_block] * num_blocks
+        all_num_frames = [self.num_frame_per_block] * 7
+        # all_num_frames[0] = 1
         if self.independent_first_frame and initial_latent is None:
             all_num_frames = [1] + all_num_frames
         num_denoising_steps = len(self.denoising_step_list)
@@ -330,12 +291,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                                  device=noise.device)
         # if exit_flags is None:
         #     exit_flags = self.generate_and_sync_list(len(all_num_frames), 2, device=noise.device)
-            # exit_flags[0] += 2
+        #     exit_flags[0] += 2
         exit_timestep = self.denoising_step_list[exit_flags[0]].item()
 
         start_gradient_frame_index = max(0, num_output_frames - 21)
 
         high_noise_timesteps = self.denoising_step_list[self.denoising_step_list >= self.boundary_timestep]
+        assert len(high_noise_timesteps) == len(self.denoising_step_list) // 2, f"Incorrect number of high noise timesteps: {len(high_noise_timesteps)}"
 
         for block_index, current_num_frames in enumerate(all_num_frames):
             noisy_input = noise[:, current_start_frame -
@@ -381,7 +343,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                             start_frame=current_start_frame).permute(
                                 0, 2, 1, 3, 4)
 
-                        if self.boundary_timestep is not None and current_timestep >= self.boundary_timestep:
+                        if self.boundary_timestep is not None and self.training_args.use_add_noise_high and current_timestep >= self.boundary_timestep:
                             denoised_pred = pred_noise_to_x_bound(
                                 pred_noise=pred_flow.flatten(0, 1),
                                 noise_input_latent=noisy_input.flatten(0, 1),
@@ -399,7 +361,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
                         next_timestep = self.denoising_step_list[index + 1]
 
-                        if self.boundary_timestep is not None and index < len(high_noise_timesteps) - 1:
+                        if self.boundary_timestep is not None and self.training_args.use_add_noise_high and index < len(high_noise_timesteps) - 1:
                             noisy_input = self.noise_scheduler.add_noise_high(
                                 denoised_pred.flatten(0, 1),
                                 torch.randn_like(denoised_pred.flatten(0, 1)),
@@ -412,8 +374,35 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                            device=noise.device,
                                            dtype=torch.long)).unflatten(
                                     0, denoised_pred.shape[:2])
-                        elif self.boundary_timestep is not None and index == len(high_noise_timesteps) - 1:
+                        elif self.boundary_timestep is not None and self.training_args.use_add_noise_high and index == len(high_noise_timesteps) - 1:
+                            assert exit_timestep < self.boundary_timestep, "exit timestep is greater than boundary timestep"
                             noisy_input = denoised_pred
+
+                            # Step 3.3: rerun with boundary timestep to update the high noise's cache
+                            # context_timestep = torch.ones_like(timestep) * self.boundary_timestep
+                            # context_denoised_pred = self.noise_scheduler.add_noise_high(
+                            #     denoised_pred.flatten(0, 1),
+                            #     torch.randn_like(denoised_pred.flatten(0, 1)),
+                            #     context_timestep,
+                            #     self.boundary_timestep * torch.ones_like(context_timestep)).unflatten(0, denoised_pred.shape[:2])
+
+                            # training_batch_bound = self._build_distill_input_kwargs(
+                            #     context_denoised_pred, context_timestep,
+                            #     training_batch.conditional_dict, training_batch)
+
+                            # current_model(
+                            #     hidden_states=training_batch_bound.
+                            #     input_kwargs['hidden_states'],
+                            #     encoder_hidden_states=training_batch_bound.
+                            #     input_kwargs['encoder_hidden_states'],
+                            #     timestep=training_batch_bound.input_kwargs['timestep'],
+                            #     encoder_hidden_states_image=training_batch_bound.
+                            #     input_kwargs.get('encoder_hidden_states_image'),
+                            #     kv_cache=kv_cache1,
+                            #     crossattn_cache=crossattn_cache,
+                            #     current_start=current_start_frame * self.frame_seq_length,
+                            #     start_frame=current_start_frame)
+
                         else:
                             noisy_input = self.noise_scheduler.add_noise(
                                 denoised_pred.flatten(0, 1),
@@ -424,7 +413,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                         dtype=torch.long)).unflatten(
                                             0, denoised_pred.shape[:2])
                 else:
-                    logger.info("Exit timestep in _generator_multi_step_simulation_forward(): %s", current_timestep)
+                    # logger.info("Exit timestep in _generator_multi_step_simulation_forward(): %s", current_timestep)
                     # Final prediction with gradient control
                     if current_start_frame < start_gradient_frame_index:
                         with torch.no_grad():
@@ -468,7 +457,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                             start_frame=current_start_frame).permute(
                                 0, 2, 1, 3, 4)
 
-                    if self.boundary_timestep is not None and self.transformer_2 is not None and current_timestep >= self.boundary_timestep:
+                    if self.boundary_timestep is not None and self.training_args.use_add_noise_high and current_timestep >= self.boundary_timestep:
                         denoised_pred = pred_noise_to_x_bound(
                             pred_noise=pred_flow.flatten(0, 1),
                             noise_input_latent=noisy_input.flatten(0, 1),
@@ -479,7 +468,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                             pred_noise=pred_flow.flatten(0, 1),
                             noise_input_latent=noisy_input.flatten(0, 1),
                             timestep=timestep,
-                            scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2])
+                            scheduler=self.noise_scheduler).unflatten(
+                                0, pred_flow.shape[:2])
                     else:
                         denoised_pred = pred_noise_to_pred_video(
                             pred_noise=pred_flow.flatten(0, 1),
@@ -487,14 +477,17 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                             timestep=timestep,
                             scheduler=self.noise_scheduler).unflatten(
                                 0, pred_flow.shape[:2])
-                        clean_denoised_pred = denoised_pred
+                        clean_denoised_pred = None
                     break
 
+            assert exit_timestep == current_timestep, "exit timestep is not equal to current timestep"
+
             # Step 3.2: record the model's output
             output[:, current_start_frame:current_start_frame +
                    current_num_frames] = denoised_pred
-            clean_output[:, current_start_frame:current_start_frame +
-                   current_num_frames] = clean_denoised_pred
+            if clean_denoised_pred is not None:
+                clean_output[:, current_start_frame:current_start_frame +
+                    current_num_frames] = clean_denoised_pred
 
             # Step 3.3: rerun with timestep zero to update the cache
             if self.boundary_timestep is not None:
@@ -503,19 +496,23 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     current_model = self.transformer_2
                 else:
                     # For high noise expert, the context timestep should be the boundary timestep
-                    context_timestep = torch.ones_like(timestep) * self.boundary_timestep
+                    # context_timestep = torch.ones_like(timestep) * self.boundary_timestep
+                    context_timestep = torch.ones_like(timestep) * self.context_noise
                     current_model = self.transformer
             else:
                 context_timestep = torch.ones_like(timestep) * self.context_noise
                 current_model = self.transformer
 
-            if self.boundary_timestep is not None and self.transformer_2 is not None and current_timestep >= self.boundary_timestep:
-                # raise ValueError("Exit timestep is greater than boundary timestep")
-                denoised_pred = self.noise_scheduler.add_noise_high(
-                    denoised_pred.flatten(0, 1),
-                    torch.randn_like(denoised_pred.flatten(0, 1)),
-                    context_timestep,
-                    self.boundary_timestep * torch.ones_like(context_timestep)).unflatten(0, denoised_pred.shape[:2])
+            if self.boundary_timestep is not None and self.training_args.use_add_noise_high and current_timestep >= self.boundary_timestep:
+                # denoised_pred = self.noise_scheduler.add_noise_high(
+                #     denoised_pred.flatten(0, 1),
+                #     torch.randn_like(denoised_pred.flatten(0, 1)),
+                #     context_timestep,
+                #     self.boundary_timestep * torch.ones_like(context_timestep)).unflatten(0, denoised_pred.shape[:2])
+                denoised_pred = self.noise_scheduler.add_noise(
+                    clean_denoised_pred.flatten(0, 1),
+                    torch.randn_like(clean_denoised_pred.flatten(0, 1)),
+                    context_timestep).unflatten(0, clean_denoised_pred.shape[:2])
             else:
                 denoised_pred = self.noise_scheduler.add_noise(
                     denoised_pred.flatten(0, 1),
@@ -527,6 +524,20 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     denoised_pred, context_timestep,
                     training_batch.conditional_dict, training_batch)
 
+                if self.boundary_timestep is not None and exit_timestep < self.boundary_timestep:
+                    self.transformer(
+                        hidden_states=training_batch_temp.
+                        input_kwargs['hidden_states'],
+                        encoder_hidden_states=training_batch_temp.
+                        input_kwargs['encoder_hidden_states'],
+                        timestep=training_batch_temp.input_kwargs['timestep'],
+                        encoder_hidden_states_image=training_batch_temp.
+                        input_kwargs.get('encoder_hidden_states_image'),
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                        start_frame=current_start_frame)
+
                 # context_timestep is 0 so we use transformer_2
                 # current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
                 current_model(
@@ -591,16 +602,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         else:
             pred_image_or_video_last_21 = pred_image_or_video
 
-        # Set up gradient mask if we generated more than minimum frames
-        if num_generated_frames != min_num_frames:
-            # Currently, we do not use gradient for the first chunk, since it contains image latents
-            gradient_mask = torch.ones_like(pred_image_or_video_last_21,
-                                            dtype=torch.bool)
-            if self.independent_first_frame:
-                gradient_mask[:, :1] = False
-            else:
-                gradient_mask[:, :self.num_frame_per_block] = False
-
         # Apply gradient masking if needed
         final_output = pred_image_or_video_last_21.to(dtype)
         if gradient_mask is not None:
@@ -630,16 +631,19 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                 min_num_frames, dtype=torch.float32, device=self.device)
 
         # Clean up caches - properly free GPU memory
-        # self._cleanup_simulation_caches(kv_cache1, crossattn_cache)
-        # gc.collect()
-        # torch.cuda.empty_cache()
+        self._cleanup_simulation_caches(kv_cache1, crossattn_cache)
+        gc.collect()
+        torch.cuda.empty_cache()
         # assert kv_cache1 is not None
         # assert crossattn_cache is not None
         # self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)
 
+        assert pred_image_or_video.shape[1] == 21, "pred_image_or_video must have 21 frames"
+
         if gradient_mask is not None:
             return final_output, exit_timestep
         else:
+            # return pred_image_or_video, exit_timestep
             return pred_image_or_video, clean_output, exit_timestep
 
     def _initialize_simulation_caches(
@@ -673,6 +677,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             max_num_frames = num_frames
         num_max_frames = max(max_num_frames, num_frames)
         kv_cache_size = num_max_frames * frame_seq_length
+        assert kv_cache_size == 1560 * 21, "kv_cache_size must be 1560 * 21"
 
         kv_cache = []
         for _ in range(num_transformer_blocks):
@@ -855,27 +860,32 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     (generator_loss / gradient_accumulation_steps).backward()
 
                 # Ensure that only one of the two transformers have received gradients
-                # if self.train_transformer_2:
-                #     # Assert that all gradients are None for transformer
-                #     raise ValueError("Transformer 2 is training")
-                #     assert all(p.grad is None for p in self.transformer.parameters())
-                #     grad_sum = 0
-                #     for n, p in self.transformer_2.named_parameters():
-                #         if p.grad is not None:
-                #             grad_sum += p.grad.sum().item()
-                #         else:
-                #             raise ValueError("Transformer 2 param %s has no gradient", n)
-                #     assert grad_sum != 0, "Transformer 2 did not receive gradients"
-                # else:
-                #     # raise ValueError("Transformer 2 is not training")
-                #     assert all(p.grad is None for p in self.transformer_2.parameters())
-                #     grad_norm = 0.0
-                #     for n, p in self.transformer.named_parameters():
-                #         if p.grad is not None:
-                #             grad_norm += p.grad.data.norm(2).item() ** 2
-                #         else:
-                #             raise ValueError("Transformer param %s has no gradient", n)
-                #     assert grad_norm > 0.0, "Transformer has 0 gradient norm"
+                if self.train_transformer_2:
+                    # Assert that all gradients are None for transformer
+                    assert all(p.grad is None for p in self.transformer.parameters())
+                    grad_sum = 0
+                    for n, p in self.transformer_2.named_parameters():
+                        if not p.requires_grad:
+                            logger.info("Transformer 2 param %s is not trainable", n)
+                            continue
+                        if p.grad is not None:
+                            grad_sum += p.grad.sum().item()
+                        else:
+                            raise ValueError("Transformer 2 param %s has no gradient", n)
+                    assert grad_sum != 0, "Transformer 2 did not receive gradients"
+                else:
+                    # raise ValueError("Transformer 2 is not training")
+                    assert all(p.grad is None for p in self.transformer_2.parameters())
+                    grad_norm = 0.0
+                    for n, p in self.transformer.named_parameters():
+                        if not p.requires_grad:
+                            logger.info("Transformer param %s is not trainable", n)
+                            continue
+                        if p.grad is not None:
+                            grad_norm += p.grad.data.norm(2).item() ** 2
+                        else:
+                            raise ValueError("Transformer param %s has no gradient", n)
+                    assert grad_norm > 0.0, "Transformer has 0 gradient norm"
 
                 total_generator_loss += generator_loss.detach().item()
                 generator_log_dict.update(gen_log_dict)
@@ -888,14 +898,32 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             if hasattr(
                     self, 'train_transformer_2'
             ) and self.train_transformer_2 and self.transformer_2 is not None:
-                self._clip_model_grad_norm_(batch_gen, self.transformer_2)
+                generator_grad_norm = self._clip_model_grad_norm_(batch_gen, self.transformer_2)
                 self.optimizer_2.step()
                 self.lr_scheduler_2.step()
             else:
-                self._clip_model_grad_norm_(batch_gen, self.transformer)
+                generator_grad_norm = self._clip_model_grad_norm_(batch_gen, self.transformer)
                 self.optimizer.step()
                 self.lr_scheduler.step()
 
+            if (self.current_trainstep >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(
+                    self.transformer, decay=self.training_args.ema_decay)
+                logger.info("Created generator EMA at step %s with decay=%s",
+                            self.current_trainstep, self.training_args.ema_decay)
+
+                # Create EMA for transformer_2 if it exists
+                if self.transformer_2 is not None and self.generator_ema_2 is None:
+                    self.generator_ema_2 = EMA_FSDP(
+                        self.transformer_2, decay=self.training_args.ema_decay)
+                    logger.info(
+                        "Created generator EMA_2 at step %s with decay=%s",
+                        self.current_trainstep, self.training_args.ema_decay)
+                
+                world_group = get_world_group()
+                world_group.barrier()
+
             if self.generator_ema is not None:
                 if hasattr(
                         self, 'train_transformer_2'
@@ -913,13 +941,18 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             world_group.all_reduce(avg_generator_loss,
                                    op=torch.distributed.ReduceOp.AVG)
             training_batch.generator_loss = avg_generator_loss.item()
-            # training_batch.fake_score_loss = 0
-            # training_batch.total_loss = training_batch.generator_loss
-            # return training_batch
+            training_batch.generator_grad_norm = generator_grad_norm
+            training_batch.critic_grad_norm = 0.0
+            training_batch.fake_score_loss = 0
+            training_batch.total_loss = training_batch.generator_loss
+            training_batch.fake_score_latent_vis_dict = {
+                "fake_score_timestep": torch.tensor(0.0, device=self.device),
+            }
+            return training_batch
         else:
             training_batch.generator_loss = 0.0
+            training_batch.generator_grad_norm = 0.0
 
-        logger.debug("Training critic at step %s", self.current_trainstep)
         self.fake_score_optimizer.zero_grad()
         if self.fake_score_transformer_2 is not None:
             self.fake_score_optimizer_2.zero_grad()
@@ -948,25 +981,25 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                      attn_metadata=batch_critic.attn_metadata):
                 (critic_loss / gradient_accumulation_steps).backward()
 
-            # if self.train_fake_score_transformer_2:
-            #     assert all(p.grad is None for p in self.fake_score_transformer.parameters())
-            #     grad_sum = 0
-            #     for n, p in self.fake_score_transformer_2.named_parameters():
-            #         if p.grad is not None:
-            #             grad_sum += p.grad.sum().item()
-            #         else:
-            #             raise ValueError("Fake score transformer 2 param %s has no gradient", n)
-            #     assert grad_sum != 0, "Fake score transformer 2 did not receive gradients"
-            # else:
-            #     raise ValueError("Fake score transformer 2 is not training")
-            #     assert all(p.grad is None for p in self.fake_score_transformer_2.parameters())
-            #     grad_sum = 0
-            #     for n, p in self.fake_score_transformer.named_parameters():
-            #         if p.grad is not None:
-            #             grad_sum += p.grad.sum().item()
-            #         else:
-            #             raise ValueError("Fake score transformer param %s has no gradient", n)
-            #     assert grad_sum != 0, "Fake score transformer did not receive gradients"
+            if self.train_fake_score_transformer_2:
+                assert all(p.grad is None for p in self.fake_score_transformer.parameters())
+                grad_sum = 0
+                for n, p in self.fake_score_transformer_2.named_parameters():
+                    if p.grad is not None:
+                        grad_sum += p.grad.sum().item()
+                    else:
+                        raise ValueError("Fake score transformer 2 param %s has no gradient", n)
+                assert grad_sum != 0, "Fake score transformer 2 did not receive gradients"
+            else:
+                if self.fake_score_transformer_2 is not None:
+                    assert all(p.grad is None for p in self.fake_score_transformer_2.parameters())
+                grad_sum = 0
+                for n, p in self.fake_score_transformer.named_parameters():
+                    if p.grad is not None:
+                        grad_sum += p.grad.sum().item()
+                    else:
+                        raise ValueError("Fake score transformer param %s has no gradient", n)
+                assert grad_sum != 0, "Fake score transformer did not receive gradients"
 
             total_critic_loss += critic_loss.detach().item()
             critic_log_dict.update(crit_log_dict)
@@ -976,12 +1009,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     batch_critic.fake_score_latent_vis_dict)
 
         if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
-            self._clip_model_grad_norm_(batch_critic,
+            critic_grad_norm = self._clip_model_grad_norm_(batch_critic,
                                         self.fake_score_transformer_2)
             self.fake_score_optimizer_2.step()
             self.fake_score_lr_scheduler_2.step()
         else:
-            self._clip_model_grad_norm_(batch_critic,
+            critic_grad_norm = self._clip_model_grad_norm_(batch_critic,
                                         self.fake_score_transformer)
             self.fake_score_optimizer.step()
             self.fake_score_lr_scheduler.step()
@@ -993,7 +1026,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         world_group.all_reduce(avg_critic_loss,
                                op=torch.distributed.ReduceOp.AVG)
         training_batch.fake_score_loss = avg_critic_loss.item()
-
+        training_batch.critic_grad_norm = critic_grad_norm
         training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
         return training_batch
 
@@ -1006,6 +1039,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
     def visualize_intermediate_latents(self, training_batch: TrainingBatch,
                                        training_args: TrainingArgs, step: int):
         """Add visualization data to wandb logging and save frames to disk."""
+        self.vae = self.vae.to(get_local_torch_device(), dtype=torch.bfloat16)
         wandb_loss_dict = {}
 
         # Debug logging
@@ -1104,6 +1138,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                         video, fps=24, format="mp4")
                     del video, latents
 
+        self.vae = self.vae.to("cpu")
+
         # Log metadata
         if hasattr(
                 training_batch,
@@ -1165,8 +1201,15 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
         step_times: deque[float] = deque(maxlen=100)
 
         self._log_training_info()
-        self._log_validation(self.transformer, self.training_args,
+        assert self.transformer.scale_shift_table.requires_grad == False
+        if self.transformer_2 is not None:
+            assert self.transformer_2.scale_shift_table.requires_grad == False
+        validation_args = deepcopy(self.training_args)
+        self._log_validation(validation_args,
                              self.init_steps)
+        assert self.transformer.scale_shift_table.requires_grad == False
+        if self.transformer_2 is not None:
+            assert self.transformer_2.scale_shift_table.requires_grad == False
         gc.collect()
         torch.cuda.empty_cache()
 
@@ -1198,20 +1241,23 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             self.current_trainstep = step
             training_batch.current_vsa_sparsity = current_vsa_sparsity
 
-            if (step >= self.training_args.ema_start_step) and \
-                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
-                self.generator_ema = EMA_FSDP(
-                    self.transformer, decay=self.training_args.ema_decay)
-                logger.info("Created generator EMA at step %s with decay=%s",
-                            step, self.training_args.ema_decay)
+            # if (step >= self.training_args.ema_start_step) and \
+            #         (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+            #     self.generator_ema = EMA_FSDP(
+            #         self.transformer, decay=self.training_args.ema_decay)
+            #     logger.info("Created generator EMA at step %s with decay=%s",
+            #                 step, self.training_args.ema_decay)
 
-                # Create EMA for transformer_2 if it exists
-                if self.transformer_2 is not None and self.generator_ema_2 is None:
-                    self.generator_ema_2 = EMA_FSDP(
-                        self.transformer_2, decay=self.training_args.ema_decay)
-                    logger.info(
-                        "Created generator EMA_2 at step %s with decay=%s",
-                        step, self.training_args.ema_decay)
+            #     # Create EMA for transformer_2 if it exists
+            #     if self.transformer_2 is not None and self.generator_ema_2 is None:
+            #         self.generator_ema_2 = EMA_FSDP(
+            #             self.transformer_2, decay=self.training_args.ema_decay)
+            #         logger.info(
+            #             "Created generator EMA_2 at step %s with decay=%s",
+            #             step, self.training_args.ema_decay)
+                
+            #     world_group = get_world_group()
+            #     world_group.barrier()
 
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 training_batch = self.train_one_step(training_batch)
@@ -1219,7 +1265,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
             total_loss = training_batch.total_loss
             generator_loss = training_batch.generator_loss
             fake_score_loss = training_batch.fake_score_loss
-            grad_norm = training_batch.grad_norm
+            generator_grad_norm = training_batch.generator_grad_norm
+            critic_grad_norm = training_batch.critic_grad_norm
 
             step_time = time.perf_counter() - start_time
             step_times.append(step_time)
@@ -1234,8 +1281,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                 f"{fake_score_loss:.4f}",
                 "step_time":
                 f"{step_time:.2f}s",
-                "grad_norm":
-                grad_norm,
+                "generator_grad_norm":
+                generator_grad_norm,
+                "critic_grad_norm":
+                critic_grad_norm,
                 "ema":
                 "✓" if (self.generator_ema is not None and self.is_ema_ready())
                 else "✗",
@@ -1259,8 +1308,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                     step_time,
                     "avg_step_time":
                     avg_step_time,
-                    "grad_norm":
-                    grad_norm,
+                    "generator_grad_norm":
+                    generator_grad_norm,
+                    "critic_grad_norm":
+                    critic_grad_norm,
                 }
                 if (step % self.dfake_gen_update_ratio == 0):
                     log_data["train_generator_loss"] = generator_loss
@@ -1340,6 +1391,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
 
                 if self.transformer:
                     self.transformer.train()
+                if getattr(self, 'transformer_2', None) is not None:
+                    self.transformer_2.train()
                 self.sp_group.barrier()
 
             if (self.training_args.weight_only_checkpointing_steps > 0
@@ -1372,11 +1425,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                                    None),
                     generator_ema_2=getattr(self, 'generator_ema_2', None))
 
-                if self.training_args.use_ema and self.is_ema_ready():
-                    self.save_ema_weights(self.training_args.output_dir, step)
+                # if self.training_args.use_ema and self.is_ema_ready():
+                #     self.save_ema_weights(self.training_args.output_dir, step)
+
+                self.sp_group.barrier()
 
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                self._log_validation(self.transformer, self.training_args, step)
+                self._log_validation(deepcopy(self.training_args), step)
                 gc.collect()
                 torch.cuda.empty_cache()
 
@@ -1412,9 +1467,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                            None),
             generator_ema_2=getattr(self, 'generator_ema_2', None))
 
-        if self.training_args.use_ema and self.is_ema_ready():
-            self.save_ema_weights(self.training_args.output_dir,
-                                  self.training_args.max_train_steps)
+        # if self.training_args.use_ema and self.is_ema_ready():
+        #     self.save_ema_weights(self.training_args.output_dir,
+        #                           self.training_args.max_train_steps)
 
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/training/training_pipeline.py

@@ -102,25 +102,21 @@ class TrainingPipeline(LoRAPipeline, ABC):
         self.transformer_2 = self.get_module("transformer_2", None)
         self.seed = training_args.seed
         self.set_schemas()
+        self.set_trainable()
 
         # Set random seeds for deterministic training
         assert self.seed is not None, "seed must be set"
         set_random_seed(self.seed)
         self.transformer.train()
         self.transformer.requires_grad_(True)
+        self.transformer.scale_shift_table.requires_grad_(False)
         if training_args.enable_gradient_checkpointing_type is not None:
             self.transformer = apply_activation_checkpointing(
                 self.transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
-            if self.transformer_2 is not None:
-                self.transformer_2 = apply_activation_checkpointing(
-                    self.transformer_2,
-                    checkpointing_type=training_args.
-                    enable_gradient_checkpointing_type)
 
         noise_scheduler = self.modules["scheduler"]
-        self.set_trainable()
         params_to_optimize = self.transformer.parameters()
         params_to_optimize = list(
             filter(lambda p: p.requires_grad, params_to_optimize))
@@ -153,6 +149,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
             # Ensure transformer_2 has trainable parameters before creating optimizer
             self.transformer_2.train()
             self.transformer_2.requires_grad_(True)
+            self.transformer_2.scale_shift_table.requires_grad_(False)
+            if training_args.enable_gradient_checkpointing_type is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
             params_to_optimize_2 = self.transformer_2.parameters()
             params_to_optimize_2 = list(
                 filter(lambda p: p.requires_grad, params_to_optimize_2))
@@ -233,18 +235,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _enable_training(self, model: torch.nn.Module,
                          optimizer: torch.optim.Optimizer) -> None:
         """Enable training mode and gradients for the specified model."""
-        for param in model.parameters():
-            param.requires_grad = True
-        model.train()
+        # for param in model.parameters():
+        #     param.requires_grad = True
+        # model.train()
         # optimizer.zero_grad()
+        return
 
     def _disable_training(self, model: torch.nn.Module,
                           optimizer: torch.optim.Optimizer) -> None:
         """Disable training mode and gradients for the specified model."""
-        for param in model.parameters():
-            param.requires_grad = False
-        model.eval()
+        # for param in model.parameters():
+        #     param.requires_grad = False
+        # model.eval()
         # optimizer.zero_grad(set_to_none=True)
+        return
 
     def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
         batch = next(self.train_loader_iter, None)  # type: ignore
@@ -714,6 +718,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
                                   validation_batch: dict[str, Any],
                                   num_inference_steps: int) -> ForwardBatch:
         sampling_param.prompt = validation_batch['prompt']
+        if 'image_path' in validation_batch:
+            sampling_param.image_path = validation_batch['image_path']
         sampling_param.height = training_args.num_height
         sampling_param.width = training_args.num_width
         sampling_param.num_inference_steps = num_inference_steps

fastvideo/training/training_utils.py

@@ -37,7 +37,7 @@ def gather_state_dict_on_cpu_rank0(
     sharded_sd = model.state_dict()
     param_requires_grad = set([
         k.replace("._checkpoint_wrapped_module.", ".")
-        for k, v in dict(model.named_parameters()).items() if v.requires_grad
+        for k, v in dict(model.named_parameters()).items()
     ])
     for param_name, param in sharded_sd.items():
         if param_name not in param_requires_grad:
@@ -255,8 +255,6 @@ def save_distillation_checkpoint(
         if generator_scheduler is not None:
             generator_states["scheduler"] = SchedulerWrapper(
                 generator_scheduler)
-        if generator_ema is not None:
-            generator_states["ema"] = generator_ema.state_dict()
 
         generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
                                          "generator")
@@ -288,8 +286,6 @@ def save_distillation_checkpoint(
             if generator_scheduler_2 is not None:
                 generator_2_states["scheduler"] = SchedulerWrapper(
                     generator_scheduler_2)
-            if generator_ema_2 is not None:
-                generator_2_states["ema"] = generator_ema_2.state_dict()
 
             generator_2_dcp_dir = os.path.join(save_dir,
                                                "distributed_checkpoint",
@@ -415,6 +411,69 @@ def save_distillation_checkpoint(
             rank,
             local_main_process_only=False)
 
+    # Persist EMA separately to avoid shape mismatches across ranks.
+    # Supports:
+    #   - mode="rank0_full": save consolidated EMA only on rank 0
+    #   - mode="local_shard": save per-rank EMA shard for each rank
+    try:
+        if generator_ema is not None and getattr(generator_ema, "mode",
+                                                 None) == "rank0_full":
+            _save_rank0_full_ema_safetensors(generator_ema,
+                                             generator_transformer, rank,
+                                             save_dir, "generator_ema")
+        elif generator_ema is not None and getattr(generator_ema, "mode",
+                                                   None) == "local_shard":
+            # Save per-rank shard
+            ema_dir_shard = os.path.join(save_dir, "ema_local_shard")
+            os.makedirs(ema_dir_shard, exist_ok=True)
+            ema_shard_path = os.path.join(ema_dir_shard, f"generator_ema_rank{rank}.pt")
+            torch.save(generator_ema.state_dict(), ema_shard_path)
+            logger.info(
+                "rank: %s, saved generator EMA shard (local_shard) to %s",
+                rank,
+                ema_shard_path,
+                local_main_process_only=False)
+
+            # Also consolidate EMA to a single full-state file on rank 0 by applying EMA to the model and gathering
+            # _consolidate_local_shard_ema_and_save_safetensors(
+            #     generator_ema, generator_transformer, rank, save_dir,
+            #     "generator_ema")
+
+    except Exception as e:
+        logger.warning("rank: %s, failed saving EMA separately: %s",
+                       rank,
+                       str(e))
+
+    try:
+        if generator_ema_2 is not None and getattr(generator_ema_2, "mode",
+                                                   None) == "rank0_full":
+            _save_rank0_full_ema_safetensors(generator_ema_2,
+                                             generator_transformer_2, rank,
+                                             save_dir, "generator_ema_2")
+        elif generator_ema_2 is not None and getattr(generator_ema_2, "mode",
+                                                     None) == "local_shard":
+            # Save per-rank shard for EMA_2
+            ema_dir_shard_2 = os.path.join(save_dir, "ema_local_shard")
+            os.makedirs(ema_dir_shard_2, exist_ok=True)
+            ema2_shard_path = os.path.join(ema_dir_shard_2, f"generator_ema_2_rank{rank}.pt")
+            torch.save(generator_ema_2.state_dict(), ema2_shard_path)
+            logger.info(
+                "rank: %s, saved generator_2 EMA shard (local_shard) to %s",
+                rank,
+                ema2_shard_path,
+                local_main_process_only=False)
+
+            # Also consolidate EMA_2 to a single full-state file on rank 0
+            # _consolidate_local_shard_ema_and_save_safetensors(
+            #     generator_ema_2, generator_transformer_2, rank, save_dir,
+            #     "generator_ema_2")
+    except Exception as e:
+        logger.warning("rank: %s, failed saving EMA_2 separately: %s",
+                       rank,
+                       str(e))
+
+    return
+
     # Save generator model weights (consolidated) for inference
     cpu_state = gather_state_dict_on_cpu_rank0(generator_transformer,
                                                device=None)
@@ -452,46 +511,46 @@ def save_distillation_checkpoint(
         logger.info("--> distillation checkpoint saved at step %s to %s", step,
                     weight_path)
 
-        # Save generator_2 model weights (consolidated) for inference (MoE support)
-        if generator_transformer_2 is not None:
-            inference_save_dir_2 = os.path.join(
-                save_dir, "generator_2_inference_transformer")
-            cpu_state_2 = gather_state_dict_on_cpu_rank0(
-                generator_transformer_2, device=None)
+    # Save generator_2 model weights (consolidated) for inference (MoE support)
+    if generator_transformer_2 is not None:
+        inference_save_dir_2 = os.path.join(
+            save_dir, "generator_2_inference_transformer")
+        cpu_state_2 = gather_state_dict_on_cpu_rank0(
+            generator_transformer_2, device=None)
 
-            if rank == 0:
-                os.makedirs(inference_save_dir_2, exist_ok=True)
-                weight_path_2 = os.path.join(
-                    inference_save_dir_2, "diffusion_pytorch_model.safetensors")
-                logger.info(
-                    "rank: %s, saving consolidated generator_2 inference checkpoint to %s",
-                    rank,
-                    weight_path_2,
-                    local_main_process_only=False)
+        if rank == 0:
+            os.makedirs(inference_save_dir_2, exist_ok=True)
+            weight_path_2 = os.path.join(
+                inference_save_dir_2, "diffusion_pytorch_model.safetensors")
+            logger.info(
+                "rank: %s, saving consolidated generator_2 inference checkpoint to %s",
+                rank,
+                weight_path_2,
+                local_main_process_only=False)
 
-                # Convert training format to diffusers format and save
-                diffusers_state_dict_2 = custom_to_hf_state_dict(
-                    cpu_state_2,
-                    generator_transformer_2.reverse_param_names_mapping)
-                save_file(diffusers_state_dict_2, weight_path_2)
+            # Convert training format to diffusers format and save
+            diffusers_state_dict_2 = custom_to_hf_state_dict(
+                cpu_state_2,
+                generator_transformer_2.reverse_param_names_mapping)
+            save_file(diffusers_state_dict_2, weight_path_2)
 
-                logger.info(
-                    "rank: %s, consolidated generator_2 inference checkpoint saved to %s",
-                    rank,
-                    weight_path_2,
-                    local_main_process_only=False)
+            logger.info(
+                "rank: %s, consolidated generator_2 inference checkpoint saved to %s",
+                rank,
+                weight_path_2,
+                local_main_process_only=False)
 
-                # Save model config
-                config_dict_2 = generator_transformer_2.hf_config
-                if "dtype" in config_dict_2:
-                    del config_dict_2["dtype"]  # TODO
-                config_path_2 = os.path.join(inference_save_dir_2,
-                                             "config.json")
-                with open(config_path_2, "w") as f:
-                    json.dump(config_dict_2, f, indent=4)
-                logger.info(
-                    "--> generator_2 distillation checkpoint saved at step %s to %s",
-                    step, weight_path_2)
+            # Save model config
+            config_dict_2 = generator_transformer_2.hf_config
+            if "dtype" in config_dict_2:
+                del config_dict_2["dtype"]  # TODO
+            config_path_2 = os.path.join(inference_save_dir_2,
+                                            "config.json")
+            with open(config_path_2, "w") as f:
+                json.dump(config_dict_2, f, indent=4)
+            logger.info(
+                "--> generator_2 distillation checkpoint saved at step %s to %s",
+                step, weight_path_2)
 
 
 def load_checkpoint(transformer,
@@ -641,18 +700,38 @@ def load_distillation_checkpoint(
         end_time - begin_time,
         local_main_process_only=False)
 
-    # Load EMA state if available and generator_ema is provided
+    # Load EMA separately if saved in rank0_full mode
     if generator_ema is not None:
         try:
-            ema_state = generator_states.get("ema")
-            if ema_state is not None:
-                generator_ema.load_state_dict(ema_state)
-                logger.info("rank: %s, generator EMA state loaded successfully",
-                            rank)
-            else:
-                logger.info("rank: %s, no EMA state found in checkpoint", rank)
+            if getattr(generator_ema, "mode", None) == "rank0_full":
+                ema_path = os.path.join(checkpoint_path, "ema",
+                                        "generator_ema.pt")
+                if rank == 0 and os.path.exists(ema_path):
+                    ema_state = torch.load(ema_path, map_location="cpu")
+                    generator_ema.load_state_dict(ema_state)
+                    logger.info(
+                        "rank: %s, generator EMA (rank0_full) loaded from %s",
+                        rank, ema_path)
+                elif rank == 0:
+                    logger.info(
+                        "rank: %s, generator EMA file not found at %s; skipping",
+                        rank, ema_path)
+            elif getattr(generator_ema, "mode", None) == "local_shard":
+                ema_path = os.path.join(checkpoint_path, "ema_local_shard",
+                                        f"generator_ema_rank{rank}.pt")
+                if os.path.exists(ema_path):
+                    ema_state = torch.load(ema_path, map_location="cpu")
+                    generator_ema.load_state_dict(ema_state)
+                    logger.info(
+                        "rank: %s, generator EMA shard (local_shard) loaded from %s",
+                        rank, ema_path)
+                else:
+                    logger.info(
+                        "rank: %s, generator EMA shard file not found at %s; skipping",
+                        rank, ema_path)
+
         except Exception as e:
-            logger.warning("rank: %s, failed to load EMA state: %s", rank,
+            logger.warning("rank: %s, failed to load generator EMA: %s", rank,
                            str(e))
 
     # Load generator_2 distributed checkpoint (MoE support)
@@ -692,26 +771,45 @@ def load_distillation_checkpoint(
                 end_time - begin_time,
                 local_main_process_only=False)
 
-            # Load EMA_2 state if available and generator_ema_2 is provided
+            # Load EMA separately if saved in rank0_full mode
             if generator_ema_2 is not None:
                 try:
-                    ema_2_state = generator_2_states.get("ema")
-                    if ema_2_state is not None:
-                        generator_ema_2.load_state_dict(ema_2_state)
-                        logger.info(
-                            "rank: %s, generator_2 EMA state loaded successfully",
-                            rank)
-                    else:
-                        logger.info(
-                            "rank: %s, no EMA_2 state found in checkpoint",
-                            rank)
+                    if getattr(generator_ema_2, "mode", None) == "rank0_full":
+                        ema_path = os.path.join(checkpoint_path, "ema",
+                                                "generator_ema_2.pt")
+                        if rank == 0 and os.path.exists(ema_path):
+                            ema_state = torch.load(ema_path, map_location="cpu")
+                            generator_ema.load_state_dict(ema_state)
+                            logger.info(
+                                "rank: %s, generator EMA 2 (rank0_full) loaded from %s",
+                                rank, ema_path)
+                        elif rank == 0:
+                            logger.info(
+                                "rank: %s, generator EMA 2 file not found at %s; skipping",
+                                rank, ema_path)
+                    elif getattr(generator_ema_2, "mode", None) == "local_shard":
+                        ema_path = os.path.join(checkpoint_path, "ema_local_shard",
+                                                f"generator_ema_2_rank{rank}.pt")
+                        if os.path.exists(ema_path):
+                            ema_state = torch.load(ema_path, map_location="cpu")
+                            generator_ema_2.load_state_dict(ema_state)
+                            logger.info(
+                                "rank: %s, generator EMA 2 shard (local_shard) loaded from %s",
+                                rank, ema_path)
+                        else:
+                            logger.info(
+                                "rank: %s, generator EMA 2 shard file not found at %s; skipping",
+                                rank, ema_path)
+
                 except Exception as e:
-                    logger.warning("rank: %s, failed to load EMA_2 state: %s",
-                                   rank, str(e))
+                    logger.warning("rank: %s, failed to load generator EMA 2: %s", rank,
+                                str(e))
         else:
             logger.info("rank: %s, generator_2 checkpoint not found, skipping",
                         rank)
 
+    # return step
+
     # Load critic distributed checkpoint
     critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
                                   "critic")
@@ -1152,7 +1250,74 @@ def custom_to_hf_state_dict(
 
     return new_state_dict
 
-# More generalized version of shift_timestep
+
+def _save_full_ema_safetensors_from_state(
+    state_dict: dict[str, Any],
+    reverse_param_names_mapping: dict[str, tuple[str, int, int]],
+    output_path: str,
+) -> None:
+    """
+    Convert a training-format state_dict to HF format and save as safetensors.
+    """
+    diffusers_state_dict = custom_to_hf_state_dict(state_dict,
+                                                   reverse_param_names_mapping)
+    save_file(diffusers_state_dict, output_path)
+
+
+def _save_rank0_full_ema_safetensors(
+    ema: "EMA_FSDP",
+    module,
+    rank: int,
+    save_dir: str,
+    base_name: str,
+) -> None:
+    if rank != 0:
+        return
+    ema_dir = os.path.join(save_dir, "ema")
+    os.makedirs(ema_dir, exist_ok=True)
+    output_path = os.path.join(ema_dir, f"{base_name}.safetensors")
+    ema_state = ema.state_dict()
+    _save_full_ema_safetensors_from_state(ema_state,
+                                          module.reverse_param_names_mapping,
+                                          output_path)
+    logger.info(
+        "rank: %s, saved %s as consolidated EMA safetensors to %s",
+        rank,
+        base_name,
+        output_path,
+        local_main_process_only=False)
+
+
+def _consolidate_local_shard_ema_and_save_safetensors(
+    ema: "EMA_FSDP",
+    module,
+    rank: int,
+    save_dir: str,
+    base_name: str,
+) -> None:
+    try:
+        # Temporarily apply EMA to the live (sharded) module and gather full CPU state on rank 0
+        with ema.apply_to_model(module):
+            cpu_state_full = gather_state_dict_on_cpu_rank0(module, device=None)
+        if rank == 0:
+            ema_dir = os.path.join(save_dir, "ema")
+            os.makedirs(ema_dir, exist_ok=True)
+            output_path = os.path.join(ema_dir, f"{base_name}.safetensors")
+            _save_full_ema_safetensors_from_state(
+                cpu_state_full, module.reverse_param_names_mapping, output_path)
+            logger.info(
+                "rank: %s, saved consolidated %s EMA (from local_shard) as safetensors to %s",
+                rank,
+                base_name,
+                output_path,
+                local_main_process_only=False)
+    except Exception as ce:
+        logger.warning(
+            "rank: %s, failed consolidating %s EMA (local_shard): %s",
+            rank,
+            base_name,
+            str(ce))
+
 def shift_timestep(timestep: torch.Tensor, shift: float,
                    min_timestep: int = 0, max_timestep: int = 1000) -> torch.Tensor:
     if shift == 1:
@@ -1161,6 +1326,29 @@ def shift_timestep(timestep: torch.Tensor, shift: float,
     denominator = 1 + (shift - 1) * t
     return min_timestep + (max_timestep - min_timestep) * (shift * t / denominator)
 
+def unshift_timestep(shifted_timestep: torch.Tensor, shift: float,
+                     min_timestep: int = 0, max_timestep: int = 1000) -> torch.Tensor:
+    """Reverse the shift_timestep operation.
+    
+    Given a timestep that has been shifted by shift_timestep, this function
+    recovers the original timestep.
+    
+    Args:
+        shifted_timestep: The shifted timestep tensor
+        shift: The shift value that was used in shift_timestep
+        min_timestep: Minimum timestep value (default: 0)
+        max_timestep: Maximum timestep value (default: 1000)
+    
+    Returns:
+        The original (unshifted) timestep tensor
+    """
+    if shift == 1:
+        return shifted_timestep
+    s = (shifted_timestep - min_timestep) / (max_timestep - min_timestep)
+    # Solve for original normalized timestep t: s = shift * t / (1 + (shift - 1) * t)
+    # Rearranging: t = s / (shift - s * (shift - 1))
+    t = s / (shift - s * (shift - 1))
+    return min_timestep + (max_timestep - min_timestep) * t
 
 # coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team.
@@ -1616,7 +1804,8 @@ def _local_numel(p: torch.Tensor) -> int:
 def count_trainable(model: torch.nn.Module) -> int:
     return sum(_local_numel(p) for p in model.parameters() if p.requires_grad)
 
-
+import gc
+from torch.distributed.tensor import DTensor
 class EMA_FSDP:
     """
     FSDP2-friendly EMA with two modes:
@@ -1676,8 +1865,11 @@ class EMA_FSDP:
         self.shadow = {}
         for name, p in module.named_parameters():
             if not p.requires_grad:
+                logger.info(f"Parameter {name} is not trainable in ema init")
                 continue
             local = self._to_local_tensor(p.detach())
+            # if "scale_shift_table" in name or "proj_out" in name:
+            #     logger.info(f"{name} in ema init is DTensor? {isinstance(p, DTensor)} with parameter shape {p.shape} and local shape {local.shape}")
             self.shadow[name] = local.clone().float().cpu()
 
     @torch.no_grad()
@@ -1698,13 +1890,24 @@ class EMA_FSDP:
         # local_shard: update local shard EMA on every rank
         for name, p in module.named_parameters():
             if not p.requires_grad:
+                logger.info(f"Parameter {name} is not trainable in ema update")
                 continue
             local = self._to_local_tensor(p.detach())
             v_cpu = local.float().cpu()
+            # if "scale_shift_table" in name or "proj_out" in name:
+            #     logger.info(f"{name} in ema update is DTensor? {isinstance(p, DTensor)} with parameter shape {p.shape} and local shape {local.shape}")
             if name not in self.shadow:
+                logger.info(f"Parameter {name} is not in shadow in ema update")
                 self.shadow[name] = v_cpu.clone()
             else:
-                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+                if self.shadow[name].shape != v_cpu.shape:
+                    logger.info(f"Parameter {name} shape mismatch in ema update. EMA shape {self.shadow[name].shape}, local shape {v_cpu.shape}")
+                    self.shadow[name] = v_cpu.clone()
+                else:
+                    self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+            del local
+        gc.collect()
+        torch.cuda.empty_cache()
 
     def state_dict(self) -> dict[str, torch.Tensor]:
         if self.mode == "rank0_full":
@@ -1746,6 +1949,7 @@ class EMA_FSDP:
             with torch.no_grad():
                 for name, p in self.module.named_parameters():
                     if not p.requires_grad:
+                        logger.info(f"Parameter {name} is not trainable in ema apply_to_model")
                         continue
                     # Save local shard
                     p_local = EMA_FSDP._to_local_tensor(p.detach())
@@ -1757,6 +1961,7 @@ class EMA_FSDP:
                     if name in self.ema.shadow:
                         ema_cpu = self.ema.shadow[name]
                         if ema_cpu.numel() != p_local.numel():
+                            logger.info(f"Parameter {name} shape mismatch in ema apply_to_model. EMA shape {ema_cpu.shape}, local shape {p_local.shape}")
                             # Shard shape mismatch (e.g., empty shard here), skip
                             continue
                         # Copy EMA shard into local param shard
@@ -1774,10 +1979,11 @@ class EMA_FSDP:
                             continue
                         saved_local = self.saved[name]
                         if saved_local.numel() != p_local.numel():
+                            logger.info(f"Parameter {name} shape mismatch in ema apply_to_model exit. EMA shape {saved_local.shape}, local shape {p_local.shape}")
                             continue
                         p_local.copy_(saved_local)
             self.saved.clear()
             return False
 
     def apply_to_model(self, module):
-        return EMA_FSDP._ApplyEMACtx(self, module)
+        return EMA_FSDP._ApplyEMACtx(self, module)